[CI failed]Revert "[Config Refactor] Derive deploy override fields from stage config" (#3287)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

.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.19.0"
+          VLLM_VERSION: "v0.20.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: true

.buildkite/pipeline.yml

@@ -1,8 +1,9 @@
 # Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
-# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
-# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
-#
-# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
+# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
+# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
+# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
+
+# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
 steps:
   - label: ":github: Resolve skip-ci & upload pipeline"
     key: upload-ci-pipeline
@@ -61,3 +62,15 @@ steps:
       - buildkite-agent pipeline upload .buildkite/test-nightly.yml
     agents:
       queue: "cpu_queue_premerge"
+
+  # L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
+  - label: "Upload Weekly Pipeline"
+    depends_on: image-build
+    key: upload-weekly-pipeline
+    if: >-
+      (build.branch == "main" && build.env("WEEKLY") == "1") ||
+      (build.branch != "main" && build.pull_request.labels includes "weekly-test")
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-weekly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/scripts/hardware_ci/run_npu_test.sh

@@ -141,5 +141,5 @@ docker run \
     "${image_name}" \
     bash -c '
     set -e
-    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
 '

.buildkite/test-amd-merge.yml

@@ -7,7 +7,9 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
 
 - label: "Diffusion Model Test"
   agent_pool: mi325_2
@@ -111,8 +113,8 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni_expansion.py -m "advanced_model" --run-level "core_model"
 
 - label: "Omni Model Test Qwen3-Omni"
   agent_pool: mi325_2

.buildkite/test-amd-ready.yaml

@@ -7,7 +7,9 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 20m pytest -vvvv -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
 
 - label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1
@@ -112,19 +114,6 @@ steps:
               pytest -s -v tests/engine/test_async_omni_engine_abort.py
       '
 
-
-# NOTE: This test is not running any thing. It is skipped and deselected.
-# Currently it is = 1 skipped, 1 deselected, 17 warnings in 0.03s ======
-# - label: "Omni Model Test Qwen2-5-Omni"
-#   agent_pool: mi325_2
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
-
 # - label: "Omni Model Test Qwen3-Omni"
 #   agent_pool: mi325_2
 #   depends_on: amd-build

.buildkite/test-amd.yaml

@@ -88,7 +88,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"

.buildkite/test-merge.yml

@@ -172,25 +172,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 30
-    depends_on: upload-merge-pipeline
-    commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Qwen3-TTS CustomVoice E2E Test"
     depends_on: upload-merge-pipeline
     commands:
@@ -235,6 +216,28 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  # - label: "MOSS-TTS-Nano E2E Test"
+  #   depends_on: upload-merge-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_LOGGING_LEVEL=DEBUG
+  #         export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+  #         pytest -s -v tests/e2e/online_serving/test_moss_tts_nano.py tests/e2e/offline_inference/test_moss_tts_nano.py -m "advanced_model" --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"
+  #         volumes:
+  #           - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Omni Model Test with H100"
     timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
@@ -363,7 +366,7 @@ steps:
           set -e
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model"
           pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
         '
     agents:

.buildkite/test-nightly.yml

@@ -10,7 +10,7 @@ steps:
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
     steps:
-      - label: ":full_moon: Omni · Function Test"
+      - label: ":full_moon: Omni · Function Test with H100"
         timeout_in_minutes: 90
         commands:
           - pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
@@ -48,6 +48,25 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Omni · Function Test with L4"
+        timeout_in_minutes: 90
+        commands:
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              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: Omni · Doc Test with L4"
         timeout_in_minutes: 90
         commands:
@@ -205,7 +224,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 omni" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and omni" -k "not test_qwen2_5" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:
@@ -278,7 +297,7 @@ steps:
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
         timeout_in_minutes: 120
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -313,10 +332,48 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Bagel Function Test with H100"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -sv tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion and H100" --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: 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: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 60
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:

.buildkite/test-weekly.yml

@@ -0,0 +1,85 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
+steps:
+  - label: "Reliability Test - qwen3-omni"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - 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: "Reliability Test - wan22"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - 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

.claude/skills/add-diffusion-model/SKILL.md

@@ -544,6 +544,15 @@ Pipelines without `SupportsModuleOffload` fall back to scanning
 well-known attribute names (`transformer`, `text_encoder`, `vae`,
 etc.), which fails for non-standard names.
 
+### Step 12: Performance Profiling
+
+After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
+
+See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.md#3-profiling-single-stage-diffusion) guide for detailed instructions on:
+1. Using the PyTorch profiler (`profiler: "torch"`) to capture detailed CPU/CUDA traces.
+2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
+3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
+
 ---
 
 ## Iterative Development Tips

.gitattributes

@@ -27,6 +27,7 @@
 *.mov binary
 *.mp4 binary
 *.mp3 binary
+*.wav binary
 *.flv binary
 *.fla binary
 *.swf binary

.github/mergify.yml

@@ -0,0 +1,223 @@
+pull_request_rules:
+- name: label-documentation
+  description: Automatically apply documentation label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^[^/]+\.md$
+      - files~=^docs/
+      - files~=^examples/
+  actions:
+    label:
+      add:
+        - documentation
+    comment:
+      message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"
+
+- name: comment-pre-commit-failure
+  description: Comment on PR when pre-commit check fails
+  conditions:
+    - status-failure=pre-commit
+    - -closed
+    - -draft
+  actions:
+    comment:
+      message: |
+        Hi @{{author}}, the pre-commit checks have failed. Please run:
+
+        ```bash
+        uv pip install pre-commit
+        pre-commit install
+        pre-commit run --all-files
+        ```
+
+        Then, commit the changes and push to your branch.
+
+        For future commits, `pre-commit` will run automatically on changed files before each commit.
+
+        > [!TIP]
+        > <details>
+        > <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
+        > <br/>
+        > <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
+        >
+        > ```bash
+        > # For mypy (substitute "3.10" with the failing version if needed)
+        > pre-commit run --hook-stage manual mypy-3.10
+        > # For markdownlint
+        > pre-commit run --hook-stage manual markdownlint
+        > ```
+        > </details>
+
+- name: comment-dco-failure
+  description: Comment on PR when DCO check fails
+  conditions:
+    - status-failure=dco
+    - -closed
+    - -draft
+  actions:
+    comment:
+      message: |
+        Hi @{{author}}, the DCO check has failed. Please click on DCO in the Checks section for instructions on how to resolve this.
+
+- name: label-ci-build
+  description: Automatically apply ci/build label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^\.github/
+      - files~=\.buildkite/
+      - files~=^cmake/
+      - files=CMakeLists.txt
+      - files~=^docker/Dockerfile
+      - files~=^requirements.*\.txt
+      - files=setup.py
+  actions:
+    label:
+      add:
+        - ci/build
+
+- name: label-frontend
+  description: Automatically apply frontend label
+  conditions:
+    - label != stale
+    - files~=^vllm_omni/entrypoints/
+  actions:
+    label:
+      add:
+        - frontend
+
+- name: label-diffusion
+  description: Automatically apply diffusion label
+  conditions:
+    - label != stale
+    - files~=^vllm_omni/diffusion/
+  actions:
+    label:
+      add:
+        - diffusion
+
+- name: label-new-model
+  description: Automatically apply new-model label
+  conditions:
+    - label != stale
+    - or:
+        - files=vllm_omni/model_executor/models/registry.py
+        - files=vllm_omni/diffusion/registry.py
+  actions:
+    label:
+      add:
+        - new-model
+
+- name: label-npu
+  description: Automatically apply npu label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/npu/
+      - title~=(?i)NPU
+      - title~=(?i)Ascend
+  actions:
+    label:
+      add:
+        - npu
+    assign:
+      users:
+        - "gcanlin"
+
+- name: label-rocm
+  description: Automatically apply rocm label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/rocm/
+      - title~=(?i)AMD
+      - title~=(?i)ROCm
+  actions:
+    label:
+      add:
+        - rocm
+    assign:
+      users:
+        - "tjtanaa"
+
+- name: label-xpu
+  description: Automatically apply xpu label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/xpu/
+      - title~=(?i)Intel
+      - title~=(?i)XPU
+  actions:
+    label:
+      add:
+        - xpu
+    assign:
+      users:
+        - "yma11"
+
+- name: label-modelrunner
+  description: Automatically apply modelrunner label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/worker/
+      - title~=(?i)ModelRunner
+      - title~=(?i)worker
+  actions:
+    label:
+      add:
+        - ModelRunner
+    assign:
+      users:
+        - "gcanlin" # To check whether need to sync to NPU
+
+- name: auto-rebase if approved, ready, and 30 commits behind main
+  conditions:
+    - base = main
+    - label=ready
+    - "#approved-reviews-by >= 1"
+    - "#commits-behind >= 30"
+    - -closed
+    - -draft
+    - -conflict
+  actions:
+    rebase: {}
+
+- name: ping author on conflicts and add 'needs-rebase' label
+  conditions:
+    - label != stale
+    - conflict
+    - -closed
+  actions:
+    label:
+      add:
+        - needs-rebase
+    comment:
+      message: |
+       This pull request has merge conflicts that must be resolved before it can be
+       merged. Please rebase the PR, @{{author}}.
+
+       https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
+
+- name: remove 'needs-rebase' label when conflict is resolved
+  conditions:
+    - -conflict
+    - -closed
+  actions:
+    label:
+      remove:
+        - needs-rebase
+
+- name: label-bug
+  description: Automatically apply bug label
+  conditions:
+    - label != stale
+    - or:
+      - title~=(?i)\bbug\b
+      - title~=(?i)\bbugfix\b
+  actions:
+    label:
+      add:
+        - bug

.gitignore

@@ -264,6 +264,8 @@ tmp_test
 vllm_omni/_version.py
 # output files
 *.wav
+# Vendored test reference audio (small, intentionally checked in)
+!tests/assets/**/*.wav
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/

benchmarks/glm_image/README.md

@@ -0,0 +1,157 @@
+# GLM-Image Benchmarks
+
+Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
+
+## Benchmarks
+
+| Benchmark | Script | Description |
+|-----------|--------|-------------|
+| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
+| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
+| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
+
+## HuggingFace Baseline
+
+Single-request sequential inference using the reference HuggingFace pipeline.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+
+## vLLM-Omni Offline
+
+Multi-GPU offline inference with pipeline parallelism and continuous batching.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--deploy-config` | - | Deploy config YAML |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
+
+### Latency Computation
+
+In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
+
+## vLLM-Omni Online Serving
+
+### Start the server
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
+    --omni --port 8091 --host 0.0.0.0 \
+    --served-model-name glm-image
+```
+
+### Run the benchmark
+
+```bash
+# T2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode t2i --num-prompts 10 --model glm-image
+
+# I2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --num-prompts 10 --model glm-image
+
+# Custom dataset
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --dataset custom \
+    --dataset-path prompts.json --num-prompts 5
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
+| `--dataset-path` | - | JSON file path (required for `custom`) |
+| `--num-prompts` | `10` | Number of benchmark requests |
+| `--max-concurrency` | `1` | Max concurrent requests |
+| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
+| `--warmup-requests` | `1` | Warmup requests before measurement |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--seed` | - | Random seed |
+| `--model` | `default` | Model name (must match `--served-model-name`) |
+| `--host` | `localhost` | Server host |
+| `--port` | `8091` | Server port |
+| `--output-file` | - | JSON output file for metrics |
+| `--num-input-images` | `1` | Number of input images for random I2I |
+
+## Dataset
+
+The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
+
+Each entry contains:
+
+- `t2i_prompt`: Text prompt for text-to-image generation
+- `i2i_prompt`: Text prompt for image-to-image editing
+- `image_url`: Source image URL for I2I (downloaded and cached on first use)
+
+Custom datasets use the same JSON format and can be provided via `--dataset-path`.
+
+## Pipeline Timings
+
+All three benchmarks report per-stage pipeline timings (in milliseconds):
+
+| Key | Description |
+|-----|-------------|
+| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
+| `stage_0_gen_ms` | AR (autoregressive) model generation time |
+| `ar2diffusion_ms` | AR output to diffusion input conversion |
+| `stage_1_gen_ms` | Diffusion model denoising time |
+| `queue_wait_ms` | Queue wait time before processing |
+
+The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
+
+## Sample Results
+
+Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
+
+| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
+|---------|------|-------------------|--------------------|
+| HuggingFace | T2I | 72.6 | 0.014 |
+| HuggingFace | I2I | 70.9 | 0.014 |
+| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
+| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
+| vLLM-Omni Online | T2I | 38.8 | 0.026 |
+| vLLM-Omni Online | I2I | 34.7 | 0.029 |

benchmarks/glm_image/benchmark_glm_image.py

@@ -0,0 +1,464 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Online serving benchmark for GLM-Image (T2I and I2I modes).
+
+Sends requests to the /v1/chat/completions endpoint and reports end-to-end
+latency, throughput, and per-stage durations (when the server is started with
+--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
+
+Supports three dataset types:
+  - prompt:   Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
+              and sends source images from image_url.
+  - random:   Generate synthetic prompts (and random images for I2I).
+  - custom:   Load from a user-specified JSON file.
+
+Usage:
+    # T2I with prompt.json (default)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --num-prompts 10
+
+    # I2I with prompt.json (downloads source images automatically)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --num-prompts 10
+
+    # Random dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --dataset random --num-prompts 20
+
+    # Custom dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --dataset custom \
+        --dataset-path my_prompts.json --num-prompts 5
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+import sys
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import numpy as np
+import requests as sync_requests
+from PIL import Image
+from tqdm.asyncio import tqdm
+
+# Import backends from the diffusion benchmark (add parent dirs to path)
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
+from backends import RequestFuncOutput
+
+BENCHMARK_DIR = Path(__file__).resolve().parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class GLMImageRequest:
+    prompt: str
+    image_path: str | None = None  # Only for I2I mode
+
+
+def download_image(url: str) -> str:
+    """Download an image to cache and return the local path."""
+    IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = IMAGE_CACHE_DIR / fname
+    if local_path.exists():
+        return str(local_path)
+    resp = sync_requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+def encode_image_as_data_url(path: str) -> str:
+    """Encode a local image file as a base64 data URL."""
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    ext = Path(path).suffix.lower()
+    mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
+    return f"data:{mime};base64,{encoded}"
+
+
+# ---------------------------------------------------------------------------
+# Datasets
+# ---------------------------------------------------------------------------
+
+
+class PromptDataset:
+    """Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
+
+    def __init__(self, args: argparse.Namespace):
+        path = _ensure_prompt_json(args.dataset_path)
+        with open(path, encoding="utf-8") as f:
+            raw = json.load(f)
+
+        prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
+        self.items: list[GLMImageRequest] = []
+
+        for entry in raw:
+            prompt = entry.get(prompt_key, "").strip()
+            if not prompt:
+                continue
+            image_path = None
+            if args.mode == "i2i":
+                url = entry.get("image_url", "")
+                if url:
+                    image_path = download_image(url)
+            self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
+
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+class RandomDataset:
+    """Generate synthetic prompts (and optional random images for I2I)."""
+
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.num_prompts = args.num_prompts
+        self._random_image_paths: list[str] | None = None
+        if args.mode == "i2i":
+            self._random_image_paths = self._generate_random_images()
+
+    def _generate_random_images(self) -> list[str]:
+        paths: list[str] = []
+        for i in range(self.args.num_input_images):
+            img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
+            path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
+            img.save(path)
+            paths.append(path)
+        return paths
+
+    def __len__(self) -> int:
+        return self.num_prompts
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        image_path = None
+        if self._random_image_paths is not None:
+            image_path = self._random_image_paths[idx % len(self._random_image_paths)]
+        return GLMImageRequest(
+            prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
+            image_path=image_path,
+        )
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return [self[i] for i in range(len(self))]
+
+
+class CustomDataset:
+    """Load from a user-specified JSON file.
+
+    Expected format:
+    [
+        {"prompt": "A cat sitting on a windowsill"},
+        {"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
+    ]
+    """
+
+    def __init__(self, args: argparse.Namespace):
+        if not args.dataset_path:
+            raise ValueError("--dataset-path is required for custom dataset")
+        with open(args.dataset_path, encoding="utf-8") as f:
+            raw = json.load(f)
+        self.items: list[GLMImageRequest] = []
+        for item in raw:
+            self.items.append(
+                GLMImageRequest(
+                    prompt=item.get("prompt", ""),
+                    image_path=item.get("image_path"),
+                )
+            )
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+# ---------------------------------------------------------------------------
+# Async request for GLM-Image (chat completions with image support)
+# ---------------------------------------------------------------------------
+
+
+async def async_glm_image_request(
+    req: GLMImageRequest,
+    api_url: str,
+    model: str,
+    session: aiohttp.ClientSession,
+    pbar: Any,
+    args: argparse.Namespace,
+) -> RequestFuncOutput:
+    """Send a single T2I or I2I request via chat completions endpoint."""
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build messages
+    if req.image_path and args.mode == "i2i":
+        data_url = encode_image_as_data_url(req.image_path)
+        content = [
+            {"type": "text", "text": req.prompt},
+            {"type": "image_url", "image_url": {"url": data_url}},
+        ]
+    else:
+        content = req.prompt
+
+    messages = [{"role": "user", "content": content}]
+
+    extra_body: dict[str, Any] = {}
+    if args.height:
+        extra_body["height"] = args.height
+    if args.width:
+        extra_body["width"] = args.width
+    if args.num_inference_steps:
+        extra_body["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        extra_body["seed"] = args.seed
+
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            resp_content = msg.get("content", [])
+                            if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
+                                first_item = resp_content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations") or {}
+                                    output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
+                except (IndexError, TypeError, AttributeError):
+                    pass
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+async def iter_requests(n: int, request_rate: float) -> Any:
+    import random as _random
+
+    for i in range(n):
+        if request_rate != float("inf") and i > 0:
+            await asyncio.sleep(_random.expovariate(request_rate))
+        yield i
+
+
+def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
+    success = [o for o in outputs if o.success]
+    errors = [o for o in outputs if not o.success]
+    latencies = [o.latency for o in success]
+    peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
+
+    stage_duration_lists: dict[str, list[float]] = {}
+    for o in success:
+        for stage, dur in (o.stage_durations or {}).items():
+            stage_duration_lists.setdefault(stage, []).append(dur)
+
+    return {
+        "duration": total_duration,
+        "completed_requests": len(success),
+        "failed_requests": len(errors),
+        "throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
+        "stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
+        "stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
+    }
+
+
+async def benchmark(args: argparse.Namespace) -> None:
+    api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
+
+    # Load dataset
+    if args.dataset == "prompt":
+        dataset = PromptDataset(args)
+    elif args.dataset == "random":
+        dataset = RandomDataset(args)
+    elif args.dataset == "custom":
+        dataset = CustomDataset(args)
+    else:
+        raise ValueError(f"Unknown dataset: {args.dataset}")
+
+    glm_requests = dataset.get_requests()
+    print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
+
+    semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
+
+    async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
+        if semaphore:
+            async with semaphore:
+                return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+        return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+
+    async with aiohttp.ClientSession() as session:
+        # Warmup
+        if args.warmup_requests and glm_requests:
+            print(f"Running {args.warmup_requests} warmup request(s)...")
+            for i in range(args.warmup_requests):
+                await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
+
+        # Main benchmark
+        pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
+        start_time = time.perf_counter()
+        tasks = []
+        async for idx in iter_requests(len(glm_requests), args.request_rate):
+            tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
+        outputs = await asyncio.gather(*tasks)
+        total_duration = time.perf_counter() - start_time
+        pbar.close()
+
+    # Metrics
+    metrics = calculate_metrics(outputs, total_duration)
+    metrics["mode"] = args.mode
+    metrics["model"] = args.model
+    metrics["dataset"] = args.dataset
+
+    print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model}")
+    print(f"{'Dataset:':<40} {args.dataset}")
+    print("-" * 50)
+    print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
+    print(f"{'Request rate:':<40} {args.request_rate}")
+    print(f"{'Max concurrency:':<40} {args.max_concurrency}")
+    print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
+    print("-" * 50)
+    print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
+    print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
+    print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
+    print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
+    print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
+
+    if metrics["peak_memory_mb_max"] > 0:
+        print("-" * 50)
+        print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
+
+    if metrics["stage_durations_mean"]:
+        print("-" * 50)
+        print("Stage Durations Mean:")
+        for stage, val in sorted(metrics["stage_durations_mean"].items()):
+            unit = "ms" if stage.endswith("_ms") else "s"
+            print(f"  {stage + ':':<38} {val:.4f} ({unit})")
+
+    print("=" * 60)
+
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
+    parser.add_argument("--dataset-path", type=str, default=None)
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--max-concurrency", type=int, default=1)
+    parser.add_argument("--request-rate", type=float, default=float("inf"))
+    parser.add_argument("--warmup-requests", type=int, default=1)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--model", type=str, default="default")
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8091)
+    parser.add_argument("--output-file", type=str, default=None)
+    parser.add_argument("--disable-tqdm", action="store_true")
+    parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
+    args = parser.parse_args()
+    asyncio.run(benchmark(args))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/huggingface/inference.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode (text-to-image, no source images needed)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (image-to-image, downloads source images)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+
+    # With custom prompt.json
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --dataset-path prompts.json --num-prompts 5
+"""
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+HEIGHT = 1024
+WIDTH = 1024
+SEED = 42
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    """Load prompts from prompt.json and prepare per-request data."""
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        if mode == "t2i":
+            prompt_key = "t2i_prompt"
+        else:
+            prompt_key = "i2i_prompt"
+
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    """Download an image to cache_dir and return the local path."""
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    from diffusers.pipelines.glm_image import GlmImagePipeline
+
+    print("=" * 60)
+    print("GLM-Image HuggingFace Baseline Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if args.mode == "i2i":
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Load pipeline
+    print(f"\nLoading pipeline from {args.model_path} ...")
+    t0 = time.perf_counter()
+    pipe = GlmImagePipeline.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+    )
+    init_time = time.perf_counter() - t0
+    print(f"Pipeline loaded in {init_time:.2f}s")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Run benchmark
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    latencies = []
+    success = 0
+    failed = 0
+
+    print(f"\nRunning {len(items)} requests sequentially...")
+    print("-" * 60)
+
+    for i, item in enumerate(items):
+        prompt = item["prompt"]
+        gen_kwargs: dict = {
+            "prompt": prompt,
+            "height": args.height,
+            "width": args.width,
+            "num_inference_steps": args.num_inference_steps,
+            "guidance_scale": args.guidance_scale,
+            "generator": generator,
+        }
+
+        if args.mode == "i2i":
+            img_path = item.get("image_path")
+            if img_path and os.path.exists(img_path):
+                gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
+            else:
+                print(f"  [{i + 1}] SKIP: no source image")
+                failed += 1
+                continue
+
+        t_start = time.perf_counter()
+        try:
+            result = pipe(**gen_kwargs)
+            image = result.images[0]
+            elapsed = time.perf_counter() - t_start
+            latencies.append(elapsed)
+            success += 1
+
+            out_path = os.path.join(args.output_dir, f"{i:04d}.png")
+            image.save(out_path)
+            print(f"  [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
+        except Exception as e:
+            elapsed = time.perf_counter() - t_start
+            failed += 1
+            print(f"  [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
+
+    # Report
+    total_gen_time = sum(latencies) if latencies else 0
+    print("\n" + "=" * 60)
+    print("HuggingFace Baseline Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{len(items)}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+    if latencies:
+        arr = np.array(latencies)
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Save metrics JSON
+    metrics = {
+        "backend": "huggingface",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+    }
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/vllm-omni/inference.py

@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+vLLM-Omni offline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (downloads source images)
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+"""
+
+import argparse
+import json
+import math
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+from PIL import Image
+from vllm import SamplingParams
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+SEED = 42
+HEIGHT = 1024
+WIDTH = 1024
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
+    factor = 32
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    # Mode-dependent totals:
+    # - t2i: small + large + EOS
+    # - i2i: large + EOS
+    if is_i2i:
+        return large_tokens + 1
+    return small_tokens + large_tokens + 1
+
+
+def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        **gen_kw,
+    }
+
+
+def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        "multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
+        **gen_kw,
+    }
+
+
+def resolve_deploy_config(args: argparse.Namespace) -> str:
+    if args.deploy_config:
+        return args.deploy_config
+    if os.path.exists(DEFAULT_DEPLOY_CONFIG):
+        return DEFAULT_DEPLOY_CONFIG
+    fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
+    if fallback.exists():
+        return str(fallback)
+    raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    is_i2i = args.mode == "i2i"
+
+    print("=" * 60)
+    print("GLM-Image vLLM-Omni Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if is_i2i:
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Init Omni
+    deploy_config = resolve_deploy_config(args)
+    print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
+    t0 = time.perf_counter()
+
+    omni = Omni(
+        model=args.model_path,
+        deploy_config=deploy_config,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        enable_ar_profiler=args.enable_ar_profiler,
+    )
+
+    init_time = time.perf_counter() - t0
+    print(f"Initialized in {init_time:.2f}s")
+
+    # Sampling params
+    max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
+    ar_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=args.seed,
+        detokenize=False,
+        extra_args={"target_h": args.height, "target_w": args.width},
+    )
+    diff_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=args.height,
+        width=args.width,
+        seed=args.seed,
+    )
+    sampling_params_list = [ar_params, diff_params]
+
+    # Build all prompts
+    gen_kw = {
+        "seed": args.seed,
+        "num_inference_steps": args.num_inference_steps,
+        "guidance_scale": args.guidance_scale,
+    }
+    all_prompts = []
+    for item in items:
+        if is_i2i:
+            img_path = item.get("image_path")
+            if not img_path or not os.path.exists(img_path):
+                continue
+            all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
+        else:
+            all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
+
+    valid = len(all_prompts)
+    print(f"Valid prompts: {valid}")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Warmup: run 1 request to prime caches, CUDA graphs, etc.
+    if all_prompts:
+        print("Running warmup request...")
+        try:
+            warmup_prompt = [all_prompts[0]]
+            omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
+            print("Warmup done.\n")
+        except Exception as e:
+            print(f"Warmup failed (continuing): {e}")
+
+    # Run
+    print(f"\nRunning {valid} requests...")
+    print("-" * 60)
+
+    latencies = []
+    all_stage_durations: list[dict[str, float]] = []
+    success = 0
+    failed = 0
+    wall_start = time.perf_counter()
+
+    try:
+        output_idx = 0
+        for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
+            if stage_outputs.final_output_type == "image":
+                request_output = stage_outputs.request_output
+                request_id = getattr(request_output, "request_id", "")
+
+                images = getattr(request_output, "images", [])
+                if not images and hasattr(request_output, "multimodal_output"):
+                    mm = request_output.multimodal_output
+                    if isinstance(mm, dict):
+                        images = mm.get("images", [])
+
+                elapsed = time.perf_counter() - wall_start
+                if images:
+                    for img in images:
+                        if isinstance(img, Image.Image):
+                            out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
+                            img.save(out_path)
+                    success += 1
+                    latencies.append(elapsed)
+                    stage_durations = getattr(stage_outputs, "stage_durations", {})
+                    if stage_durations:
+                        all_stage_durations.append(stage_durations)
+                    # Show wall-clock elapsed and pipeline breakdown if available
+                    preprocess_str = ""
+                    if "preprocess_ms" in stage_durations:
+                        preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
+                    print(f"  [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
+                    output_idx += 1
+                else:
+                    failed += 1
+    except Exception as e:
+        print(f"Error: {e}")
+        failed = valid - success
+
+    total_gen_time = time.perf_counter() - wall_start
+
+    # Diff stage_0_gen_ms with previous request to remove accumulated wait time.
+    # stage_0_gen_ms is measured from submit_ts (same for all requests submitted
+    # at once), so it accumulates queue/scheduling overhead across requests.
+    # Other stages and pipeline timings are per-request already.
+    _TIMING_ORDER = [
+        "preprocess_ms",
+        "stage_0_gen_ms",
+        "ar2diffusion_ms",
+        "stage_1_gen_ms",
+        "queue_wait_ms",
+    ]
+
+    per_request_actual: list[dict[str, float]] = []
+    prev_stage_0_ms = 0.0
+    for sd in all_stage_durations:
+        actual = dict(sd)
+        s0 = sd.get("stage_0_gen_ms", 0.0)
+        actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
+        prev_stage_0_ms = s0
+        per_request_actual.append(actual)
+
+    per_request_e2e_ms: list[float] = []
+    for actual in per_request_actual:
+        e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
+        if e2e_ms > 0:
+            per_request_e2e_ms.append(e2e_ms)
+
+    # Report
+    print("\n" + "=" * 60)
+    print("vLLM-Omni Benchmark Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{valid}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
+        print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
+        print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
+        print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
+        print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
+        print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
+        print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
+        print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
+
+    if per_request_actual:
+        print("-" * 50)
+        print("Pipeline Timings Mean:")
+        for key in _TIMING_ORDER:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+        # Show any extra keys not in the ordered list
+        ordered_set = set(_TIMING_ORDER)
+        extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
+        for key in extra_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Metrics JSON
+    metrics = {
+        "backend": "vllm-omni",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+    }
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        metrics["latency_mean"] = float(per_request_s.mean())
+        metrics["latency_median"] = float(np.median(per_request_s))
+        metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        metrics["latency_mean"] = float(per_request.mean())
+        metrics["latency_median"] = float(np.median(per_request))
+        metrics["latency_p95"] = float(np.percentile(per_request, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request, 99))
+    else:
+        metrics["latency_mean"] = 0
+        metrics["latency_median"] = 0
+        metrics["latency_p95"] = 0
+        metrics["latency_p99"] = 0
+    if per_request_actual:
+        all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
+        stage_metrics = {}
+        for key in all_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            stage_metrics[key] = {
+                "mean": float(np.mean(vals)),
+                "median": float(np.median(vals)),
+                "p95": float(np.percentile(vals, 95)),
+            }
+        metrics["stage_durations"] = stage_metrics
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+    omni.close()
+    print("Done!")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler for stage-level timing",
+    )
+    parser.add_argument(
+        "--enable-ar-profiler",
+        action="store_true",
+        help="Enable AR stage profiler to include AR timing in stage_durations",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable detailed per-request pipeline stats logging",
+    )
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.19.0
+ARG VLLM_BASE_TAG=v0.20.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.19.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.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

@@ -74,7 +74,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.19.0
+ARG VLLM_VERSION=v0.20.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -10,6 +10,7 @@ nav:
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
+      - Audio Generation: serving/audio_generate_api.md
       - Streaming Video Input: serving/video_stream_api.md
   - Examples:
     - examples/README.md
@@ -22,6 +23,7 @@ nav:
       - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
       - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
@@ -61,11 +63,6 @@ nav:
       - Cache Acceleration:
         - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
         - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
-      - Quantization:
-        - Overview: user_guide/diffusion/quantization/overview.md
-        - FP8: user_guide/diffusion/quantization/fp8.md
-        - Int8: user_guide/diffusion/quantization/int8.md
-        - GGUF: user_guide/diffusion/quantization/gguf.md
       - Attention Backends: user_guide/diffusion/attention_backends.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Parallelism:
@@ -80,6 +77,14 @@ nav:
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md
       - Step Execution: user_guide/diffusion/step_execution.md
+    - Quantization:
+      - Overview: user_guide/quantization/overview.md
+      - Online Quantization: user_guide/quantization/online.md
+      - FP8 W8A8: user_guide/quantization/fp8.md
+      - Int8 W8A8: user_guide/quantization/int8.md
+      - GGUF: user_guide/quantization/gguf.md
+      - AutoRound: user_guide/quantization/autoround.md
+      - msModelSlim: user_guide/quantization/msmodelslim.md
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
@@ -91,7 +96,6 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
-    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md

docs/cli/serve.md

@@ -12,9 +12,11 @@ The stage-based CLI is designed for deployments that require launching each pipe
   `vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
 
 Example: Initializing Stage 0 (Orchestrator and API Server):
+The commands below show a common device mapping where Stage 0 uses GPU 0 and
+worker stages use GPU 1 via `CUDA_VISIBLE_DEVICES`.
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -24,7 +26,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 Example: Initializing a Headless Worker Stage (Stage 1):
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -46,7 +48,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

docs/configuration/stage_configs.md

@@ -102,15 +102,17 @@ The stage-based CLI paradigm facilitates the execution of discrete pipeline stag
 
 For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
 does **not** necessitate the explicit definition of `--deploy-config`:
+the example below uses `CUDA_VISIBLE_DEVICES=0` for Stage 0 and
+`CUDA_VISIBLE_DEVICES=1` for Stage 1.
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
     --omni-master-port 26000
 
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -132,7 +134,7 @@ Conversely, in the context of the **stage-based CLI** paradigm, given that each
 can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \
@@ -189,7 +191,7 @@ Within the stage-based CLI paradigm, equivalent configuration parameters can inh
 as command-line arguments to the designated single-stage process instantiation:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 0 \
     --max-num-seqs 8 \
     --omni-master-address 127.0.0.1 \

docs/contributing/ci/CI_5levels.md

@@ -108,7 +108,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         /tests/dfx/stability/tests/test_qwen3_omni.json<br>
         /tests/dfx/stability/tests/test_wan22.json<br>
         <strong>Reliability:</strong><br>
-        tests/e2e/reliability/test_{model_name}.py
+        tests/dfx/reliability/test_reliability_{model_key}.py<br>
+        (e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
       </td>
       <td>
         <a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
@@ -576,12 +577,17 @@ L5 level testing focuses on the performance of model services under ***long-runn
 ### 4.2 Testing Content and Scope
 
 -   ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
--   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
+-   ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
 
 ### 4.3 Test Directory and Execution Files
 
 -   ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
--   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
+-   ***Reliability Test Suite*** (`tests/dfx/reliability/`):
+    -   `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
+    -   `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
+    -   `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
+    -   `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
+    -   `README.md` — Short local run commands for this directory.
 
 ### 4.4 Execution Method and Example
 
@@ -662,8 +668,57 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
 
 </details>
 
+<details>
+<summary> Reliability test suite (<code>tests/dfx/reliability</code>)</summary>
+
+#### Purpose and relationship to stability
+
+Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load).
+
+#### Directory layout
+
+| Path | Responsibility |
+| ---- | -------------- |
+| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
+| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
+| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
+| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
+| `README.md` | Minimal run / collect examples. |
+
+#### Parametrization and markers
+
+- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
+- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
+
+#### CI trigger
+
+Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
+
+```bash
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+```
+
+#### Local commands
+
+```bash
+pytest --collect-only tests/dfx/reliability
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+```
+
+#### Adding a new model suite
+
+1. Add `test_reliability_<model>.py` under `tests/dfx/reliability/`.
+2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files).
+3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML.
+4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/test-weekly.yml`** when the suite should run in weekly L5.
+
+</details>
+
 -   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
+    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
 
 ## Summary
 

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -1,6 +1,6 @@
 **Scope**
 
-For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
 
 **Test Case Design**
 

docs/contributing/model/adding_quantization_model.md

@@ -1,478 +0,0 @@
-# Support FP8 Quantization
-
-This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
-
----
-
-## Table of Contents
-
-- [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
-- [Testing](#testing)
-- [Troubleshooting](#troubleshooting)
-- [Reference Implementations](#reference-implementations)
-- [Summary](#summary)
-
----
-
-## Overview
-
-### What is FP8 Quantization?
-
-FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
-
-A diffusion pipeline typically has three quantizable components:
-
-| Component | Layer Types | Quantization Approach |
-|-----------|-----------|----------------------|
-| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
-| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
-| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
-
-### Architecture
-
-The quantization system has two mechanisms:
-
-**1. vLLM Quantized Linear Layers (for DiT)**
-
-vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
-
-Key classes:
-
-| Class | Location | Purpose |
-|-------|----------|---------|
-| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
-| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
-
-**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
-
-For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
-
-```
-At rest:     weight.data = fp8_tensor (half memory)
-Pre-hook:    weight.data = fp8_weight.to(bf16) * scale
-Forward:     Normal computation in BF16
-Post-hook:   weight.data = fp8_weight (back to FP8)
-```
-
-Key function:
-
-| Function | Location | Purpose |
-|----------|----------|---------|
-| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
-
----
-
-## Step-by-Step Implementation
-
-### Step 1: Add `quant_config` to the Transformer
-
-Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
-
-#### 1.1: Accept `quant_config` in Constructor
-
-```python
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-class YourTransformer2DModel(nn.Module):
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        quant_config: QuantizationConfig | None = None,
-        # ... other params
-    ):
-        self.quant_config = quant_config
-        # Pass to sub-modules
-        self.blocks = nn.ModuleList([
-            YourTransformerBlock(quant_config=quant_config, ...)
-            for _ in range(num_layers)
-        ])
-```
-
-#### 1.2: Pass `quant_config` to vLLM Linear Layers
-
-Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
-
-```python
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-
-class YourAttentionBlock(nn.Module):
-    def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
-        # QKV projection
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=hidden_size,
-            head_size=hidden_size // num_heads,
-            total_num_heads=num_heads,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_qkv",
-        )
-        # Output projection
-        self.to_out = RowParallelLinear(
-            input_size=hidden_size,
-            output_size=hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_out",
-        )
-
-class YourFeedForward(nn.Module):
-    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
-        self.gate_up_proj = ColumnParallelLinear(
-            input_size=hidden_size,
-            output_size=intermediate_size * 2,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            input_size=intermediate_size,
-            output_size=hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-        )
-```
-
-!!! important
-    Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
-
-#### 1.3: Identify Sensitive Layers
-
-Some layers degrade output quality when quantized. Common sensitive layers:
-
-| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
-|-------|-------------|---------------|----------------------|
-| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
-| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
-| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
-| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
-| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
-
-**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
-
-| Model | Method | All Layers | With `ignored_layers` | Recommendation |
-|-------|--------|-----------|----------------------|----------------|
-| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
-| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
-| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
-
-To identify sensitive layers for a new model:
-
-1. Run inference with all layers quantized
-2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
-3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
-4. Document the recommended `ignored_layers` in the model's supported models table
-
-### Step 2: Wire Up Quantization in the Pipeline
-
-#### 2.1: Create `quant_config` and Pass to Transformer
-
-In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
-
-```python
-from vllm_omni.quantization import build_quant_config
-
-class YourPipeline(nn.Module):
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-        # ... load scheduler, text_encoder, vae ...
-
-        # Build quant config from OmniDiffusionConfig
-        quant_config = build_quant_config(od_config.quantization_config)
-
-        # Pass to transformer
-        self.transformer = YourTransformer2DModel(
-            od_config=od_config,
-            quant_config=quant_config,
-            # ... other kwargs
-        )
-```
-
-#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
-
-For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
-
-!!! note
-    The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
-
-```python
-from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
-
-class YourPipeline(nn.Module):
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-
-        self.text_encoder = SomeTextEncoder.from_pretrained(...)
-        self.vae = SomeVAE.from_pretrained(...).to(self.device)
-
-        # Apply FP8 weight storage to text encoder and VAE
-        if (
-            od_config.quantization_config is not None
-            and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
-        ):
-            apply_fp8_weight_storage(self.vae)
-            apply_fp8_weight_storage(self.text_encoder)
-
-        # ... rest of init
-```
-
-`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
-
-#### 2.3: Fix `load_weights` to Include Pre-loaded Components
-
-Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
-
-```python
-def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-    loader = AutoWeightsLoader(self)
-    loaded_weights = loader.load_weights(weights)
-    # Mark pre-loaded component weights as loaded
-    loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
-    loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
-    return loaded_weights
-```
-
-### Step 3: Add to Supported Models Documentation
-
-Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
-
-```markdown
-| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
-```
-
----
-
-## Testing
-
-### Quick Verification
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Test with FP8 enabled
-omni = Omni(model="your-model", quantization="fp8")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Quality Gate Test (LPIPS)
-
-We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
-It generates outputs with both BF16 and your quantized method using the same seed, computes
-[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
-exceeds a threshold**. This is the recommended way to validate that a quantization method does not
-unacceptably degrade output quality.
-
-**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
-
-```python
-QualityTestConfig(
-    id="int8_z_image",              # pytest ID
-    model="Tongyi-MAI/Z-Image-Turbo",
-    quantization="int8",            # your quantization method
-    task="t2i",                     # "t2i" or "t2v"
-    prompt="a cup of coffee on a wooden table, morning light",
-    max_lpips=0.15,                 # fail threshold (see guide below)
-    num_inference_steps=20,
-),
-```
-
-**Step 2: Run the test:**
-
-```bash
-# Install dependency
-pip install lpips
-
-# Run all quality tests
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
-
-# Run only your method
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
-
-# Run a specific model
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
-```
-
-**Step 3: Read the output:**
-
-```
-Quantization Quality: int8_z_image
-  Model:         Tongyi-MAI/Z-Image-Turbo
-  Method:        int8
-  LPIPS:         0.0290  (threshold: 0.15)
-  BF16 memory:   19.15 GiB
-  Quant memory:  14.80 GiB  (23% reduction)
-  Result:        PASS
-```
-
-**LPIPS threshold guide:**
-
-| LPIPS Range | Meaning | Recommendation |
-|-------------|---------|----------------|
-| < 0.01 | Imperceptible | Ideal — no visible difference |
-| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
-| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
-| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
-| > 0.30 | Significant degradation | Quantization method may not suit this model |
-
-!!! tip
-    If your test fails, don't raise the threshold — instead, identify sensitive layers and add
-    them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
-
-### Quality Benchmark (Detailed)
-
-For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
-standalone benchmark script:
-
-```bash
-python benchmarks/diffusion/quantization_quality.py \
-    --model Tongyi-MAI/Z-Image-Turbo \
-    --quantization fp8 \
-    --prompts \
-        "a cup of coffee on the table" \
-        "an aerial view of a coral reef" \
-    --height 1024 --width 1024 \
-    --num-inference-steps 50 --seed 42
-```
-
-This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
-summary table ready to paste into your PR.
-
-### Identifying Sensitive Layers with `ignored_layers`
-
-Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
-in BF16 while quantizing everything else. This is critical for achieving good quality with
-aggressive quantization methods (Int8, NVFP4).
-
-**Why some layers are sensitive:**
-
-| Layer Type | Why Sensitive | Typical Impact |
-|------------|--------------|----------------|
-| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
-| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
-| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
-| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
-| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
-
-**How to find sensitive layers for your model:**
-
-1. Run the quality gate test with all layers quantized:
-    ```bash
-    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
-    ```
-
-2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
-    ```python
-    QualityTestConfig(
-        id="fp8_your_model_skip_mlp",
-        model="your-org/Your-Model",
-        quantization="fp8",
-        task="t2i",
-        prompt="a cup of coffee on a wooden table",
-        max_lpips=0.05,
-    ),
-    ```
-
-    And in the test, modify the Omni call:
-    ```python
-    omni_qt = Omni(
-        model=config.model,
-        quantization_config={
-            "method": config.quantization,
-            "ignored_layers": ["img_mlp"],
-        },
-    )
-    ```
-
-3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
-
-4. Document the recommended `ignored_layers` in the supported models table.
-
-**Real-world examples:**
-
-| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
-|-------|--------|-------------------|----------------------|---------------------|
-| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
-| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
-| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
-
-### What to Verify (Checklist)
-
-1. **Model loads without errors** — Check logs for quantization initialization messages
-2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
-3. **Memory reduction** — Expect ~30-50% savings on quantized components
-4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
-5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
-
----
-
-## Troubleshooting
-
-### Issue: Output Quality Degradation
-
-**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
-
-**Causes & Solutions:**
-
-- **Sensitive layers being quantized:**
-
-    **Solution:** Identify and skip sensitive layers:
-    ```python
-    omni = Omni(
-        model="your-model",
-        quantization_config={
-            "method": "fp8",
-            "ignored_layers": ["img_mlp"],
-        },
-    )
-    ```
-
-- **All layers sensitive (rare):**
-
-    **Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
-
-### Issue: Missing Weights Warning
-
-**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
-
-**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
-
-### Issue: CUDA Out of Memory
-
-**Symptoms:** OOM errors during model loading despite FP8 being enabled.
-
-**Causes & Solutions:**
-
-- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
-
-    FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
-
----
-
-## Reference Implementations
-
-| Model | Pipeline File | Transformer File | Notes |
-|-------|-------------|-----------------|-------|
-| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
-| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
-
-All files are under `vllm_omni/diffusion/`.
-
----
-
-## Summary
-
-Adding FP8 quantization support to a new model:
-
-1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
-2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
-3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
-4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
-5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
-6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior

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

@@ -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.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
 
 ```bash
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto
 ```
 
-The 0.19.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.20.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.19.0
+git checkout v0.20.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
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -80,7 +80,7 @@ uv pip install --editable .
 Install vllm with command below (If you already have PyTorch).
 ```bash
 python use_existing_torch.py
-uv pip install -r requirements/build.txt
+uv pip install -r requirements/build/cuda.txt
 uv pip install --no-build-isolation --editable .
 ```
 </details>

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.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```
 
 #### 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.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
 
 ```bash
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```
 
-The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. 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.19.0
+git checkout v0.20.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.19.0 \
+  vllm/vllm-omni-rocm:v0.20.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.19.0
+  vllm/vllm-omni-rocm:v0.20.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.19.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -31,6 +31,12 @@ uv pip install -e .
 
 For additional installation methods — please see the [installation guide](installation/README.md).
 
+
+!!! 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.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+
 ## Offline Inference
 
 Text-to-image generation quickstart with vLLM-Omni:

docs/models/supported_models.md

@@ -20,6 +20,7 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/serving/audio_generate_api.md

@@ -0,0 +1,338 @@
+# Audio Generate API
+
+vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
+
+Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
+
+Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### Generate Audio
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0,
+    },
+    timeout=300.0,
+)
+
+with open("cat.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+
+#### Diffusion Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with the appropriate `Content-Type` header:
+
+| `response_format` | Content-Type |
+|--------------------|--------------|
+| `wav` | `audio/wav` |
+| `mp3` | `audio/mpeg` |
+| `flac` | `audio/flac` |
+| `pcm` | `audio/pcm` |
+| `aac` | `audio/aac` |
+| `opus` | `audio/opus` |
+
+## Examples
+
+### Basic Generation
+
+Generate audio with only a text prompt (model defaults for all other parameters):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves crashing on a beach"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+Specify an explicit audio length in seconds:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+Set a `seed` to get deterministic results across runs:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Full Control
+
+Combine all parameters for precise control over generation:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+### Python Client
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "Thunder and rain",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42,
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("thunder.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Parameter Tuning Guide
+
+### `guidance_scale`
+
+Controls how closely the generated audio follows the text prompt.
+
+| Range | Behaviour |
+|-------|-----------|
+| 3 - 5 | More creative / varied output |
+| 7 (default) | Balanced adherence |
+| 10+ | Strict adherence to the prompt |
+
+### `num_inference_steps`
+
+Controls the number of denoising steps in the diffusion process.
+
+| Steps | Quality | Speed | Use Case |
+|-------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 | Very Good | Medium | General purpose |
+| 150+ | Excellent | Slow | Final / critical audio |
+
+### `audio_length`
+
+Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
+
+### `negative_prompt`
+
+Describes characteristics to avoid. Common negative prompts include:
+
+- `"Low quality, distorted, noisy"`
+- `"Silence, static"`
+- `"Music"` (when generating sound effects only)
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid or missing parameters:
+
+```json
+{
+    "error": {
+        "message": "Audio generation model did not produce audio output.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model mismatch:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+### 422 Unprocessable Entity
+
+Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
+
+```json
+{
+    "detail": [
+        {
+            "type": "literal_error",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            ...
+        }
+    ]
+}
+```
+
+## Troubleshooting
+
+### "Audio generation model did not produce audio output"
+
+The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
+
+### Server Not Responding
+
+```bash
+# Check if the server is healthy
+curl http://localhost:8091/health
+```
+
+### Audio Quality Issues
+
+- Increase `num_inference_steps` (e.g. 150).
+- Add a negative prompt: `"Low quality, distorted, noisy"`.
+- Increase `guidance_scale` for stronger prompt adherence.
+
+### Generation Timeout
+
+- Reduce `num_inference_steps`.
+- Reduce `audio_length`.
+- Check GPU memory with `nvidia-smi`.
+
+### Out of Memory
+
+- Lower `--gpu-memory-utilization` (e.g. 0.8).
+- Reduce `audio_length`.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni \
+    --uvicorn-log-level debug
+```

docs/user_guide/diffusion/quantization/autoround.md

@@ -1,93 +0,0 @@
-# AutoRound Quantization
-
-## Overview
-
-[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
-
-The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W4A16 | 4 | ✅ Supported |
-| W8A16 | 8 | Planned |
-
-W4A16 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Configuration
-
-1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=28),
-)
-outputs[0].save_images("output.png")
-```
-
-2. **CLI**: pass the quantized model path directly.
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
-  --prompt "A cat sitting on a windowsill" \
-  --num-inference-steps 28 \
-  --output outputs/flux_w4a16.png
-```
-
-No `--quantization` flag is needed — the quantization method is read from the checkpoint.
-
-## How It Works
-
-The checkpoint's `config.json` contains:
-
-```json
-{
-  "quantization_config": {
-    "quant_method": "auto-round",
-    "bits": 4,
-    "group_size": 128,
-    "sym": true,
-    "packing_format": "auto_round:auto_gptq",
-    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
-  }
-}
-```
-
-At load time:
-
-1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
-2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
-3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
-
-## Supported Models
-
-| Model | HF Checkpoint | Scheme | Group Size | Backend |
-|-------|--------------|--------|------------|---------|
-| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
-| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
-| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
-
-## Creating a Quantized Checkpoint
-
-Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
-
-```bash
-# W4A16 (4-bit weight, 16-bit activation)
-auto-round \
-    --model black-forest-labs/FLUX.1-dev \
-    --scheme W4A16 \
-    --batch_size 1 \
-    --disable_opt_rtn \
-    --dataset coco2014 \
-    --iters 0
-```
-
-The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.

docs/user_guide/diffusion/quantization/fp8.md

@@ -1,82 +0,0 @@
-# FP8 Quantization
-
-## Overview
-
-FP8 quantization converts BF16/FP16 weights to FP8 at model load time. No calibration or pre-quantized checkpoint needed.
-
-Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16. See the [per-model table](#supported-models) for which case applies.
-
-Common sensitive layers in DiT-based diffusion models include **image-stream MLPs** (`img_mlp`). These are particularly vulnerable to FP8 precision loss because they process denoising latents whose dynamic range shifts significantly across timesteps, and unlike attention projections (which benefit from QK-Norm stabilization), MLPs have no built-in normalization to absorb quantization error. In deep architectures (e.g., 60+ residual blocks), small per-layer errors compound and degrade output quality. Other layers such as **attention projections** (`to_qkv`, `to_out`) and **text-stream MLPs** (`txt_mlp`) are generally more robust due to normalization or more stable input statistics.
-
-## Configuration
-
-1. **Python API**: set `quantization="fp8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# All layers quantized
-omni = Omni(model="<your-model>", quantization="fp8")
-
-# Skip sensitive layers
-omni = Omni(
-    model="<your-model>",
-    quantization_config={
-        "method": "fp8",
-        "ignored_layers": ["<layer-name>"],
-    },
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-2. **CLI**: pass `--quantization fp8` and optionally `--ignored-layers`.
-
-```bash
-# All layers
-python text_to_image.py --model <your-model> --quantization fp8
-
-# Skip sensitive layers
-python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
-
-# Online serving
-vllm serve <your-model> --omni --quantization fp8
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `method` | str | — | Quantization method (`"fp8"`) |
-| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
-| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) or `"static"` |
-| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
-
-The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
-| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
-| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
-| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
-| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
-| GLM-Image | `zai-org/GLM-Image` | All layers | None |
-| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
-
-## Combining with Other Features
-
-FP8 quantization can be combined with cache acceleration:
-
-```python
-omni = Omni(
-    model="<your-model>",
-    quantization="fp8",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2},
-)
-```

docs/user_guide/diffusion/quantization/gguf.md

@@ -1,227 +0,0 @@
-# GGUF Quantization
-
-## Goals
-1. Reuse vLLM quantization configs and weight loaders as much as possible.
-2. Add native GGUF support to diffusion transformers without changing model definitions.
-3. Keep user-facing knobs minimal and consistent across offline and online flows.
-
-## Scope
-1. Models: Z-Image, Flux2-klein, and Qwen-Image family.
-2. Components: diffusion transformer weights, loader paths, and quantization configs.
-3. Modes: native GGUF (transformer-only weights).
-
-## Architecture Overview
-1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
-2. Diffusion quantization wrapper (`DiffusionGGUFConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
-3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
-4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
-5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
-
-## Call Chain (Offline)
-```
-CLI (examples/offline_inference/text_to_image/text_to_image.py)
-  |
-  v
-Omni (vllm_omni/entrypoints/__init__.py)
-  |
-  v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
-  |
-  v
-DiffusionWorker
-  |
-  v
-DiffusionModelRunner
-  |
-  v
-DiffusersPipelineLoader
-  |
-  v
-Pipeline.forward (Flux2/Qwen/Z-Image)
-  |
-  v
-DiffusionEngine
-  |
-  v
-OmniRequestOutput
-  |
-  v
-Client (saved PNG)
-```
-
-## Call Chain (Online)
-```
-Client
-  |
-  | POST /v1/images/generations
-  v
-APIServer (vllm_omni/entrypoints/openai/api_server.py)
-  |
-  v
-_generate_with_async_omni
-  |
-  v
-AsyncOmni
-  |
-  v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
-  |
-  v
-OmniRequestOutput
-  |
-  v
-encode_image_base64
-  |
-  v
-ImageGenerationResponse
-  |
-  v
-Client
-```
-
-## Call Chain (GGUF Operator Path)
-```
-Pipeline.forward (Flux2/Qwen/Z-Image)
-  |
-  v
-Transformer blocks
-  |
-  v
-QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
-  |
-  v
-LinearBase.forward
-  |
-  v
-QuantMethod.apply (GGUFLinearMethod.apply)
-  |
-  v
-fused_mul_mat_gguf
-  |
-  v
-_fused_mul_mat_gguf (custom op)
-  |
-  v
-ops.ggml_dequantize
-  |
-  v
-x @ weight.T
-```
-
-## GGUF Weight Loading Path (Transformer-Only)
-1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
-2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
-3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
-4. Name mapping is applied per-architecture (Z-Image, Flux2Klein, Qwen-Image).
-5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
-
-## GGUF Adapter Design
-1. `GGUFAdapter` is an abstract base class for model-specific adapters.
-2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
-3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
-4. `QwenImageGGUFAdapter` implements Qwen-Image fused attention projection remapping.
-5. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
-
-Adapter paths:
-- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
-- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
-- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
-- Qwen-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py`
-
-## User Usage (Offline)
-
-### Baseline BF16
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
-  --height 768 \
-  --width 1360 \
-  --seed 42 \
-  --cfg-scale 4.0 \
-  --num-images-per-prompt 1 \
-  --num-inference-steps 4 \
-  --output outputs/flux2_klein_4b.png
-```
-
-### Native GGUF (Transformer Only)
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
-  --quantization gguf \
-  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
-  --height 768 \
-  --width 1360 \
-  --seed 42 \
-  --cfg-scale 4.0 \
-  --num-images-per-prompt 1 \
-  --num-inference-steps 4 \
-  --output outputs/flux2_klein_4b_gguf.png
-```
-
-Notes for GGUF:
-1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
-2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
-
-## User Usage (Online)
-
-### Start Server (Native GGUF via CLI)
-```bash
-vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --omni \
-  --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
-```
-
-### Online Request (Images API)
-```bash
-curl -X POST http://localhost:8000/v1/images/generations \
-  -H "Content-Type: application/json" \
-  -d '{
-    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
-    "size": "1024x1024",
-    "seed": 42,
-    "num_inference_steps": 4
-  }'
-```
-
-## Qwen-Image Example
-
-### Offline
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model Qwen/Qwen-Image \
-  --gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
-  --quantization gguf \
-  --prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
-  --height 1024 \
-  --width 1024 \
-  --seed 42 \
-  --num_inference_steps 20 \
-  --output outputs/qwen_image_q4km.png
-```
-
-### Online
-```bash
-vllm serve Qwen/Qwen-Image \
-  --omni \
-  --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
-```
-
-Notes:
-1. The GGUF repo provides transformer weights only; the base model repo still provides tokenizer, text encoder, scheduler, and VAE.
-2. The same adapter path covers `QwenImagePipeline`, `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, and `QwenImageLayeredPipeline`.

docs/user_guide/diffusion/quantization/int8.md

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

docs/user_guide/diffusion/quantization/msmodelslim.md

@@ -1,56 +0,0 @@
-# msModelSlim Quantization
-
-## Overview
-
-[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend-friendly compression tool focused on acceleration, using compression techniques, and built for Ascend hardware. It includes a series of inference optimization technologies such as quantization and compression, aiming to accelerate large language dense models, MoE models, multimodal understanding models, multimodal generation models, etc.
-
-Once you have a quantized model which is generated by **msModelSlim**, you can use vLLM Omni for inference by specifying the --quantization ascend parameter to enable quantization features.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W8A8 | 8 | ✅ Supported |
-| W4A4 | 4 | Planned |
-
-W8A8 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Model Quantization
-
-The following example shows how to generate W8A8 quantized weights for the [Wan2_2 model](https://gitcode.com/Ascend/msmodelslim/blob/master/example/multimodal_sd/Wan2_2/README.md).
-
-**Quantization Script:**
-
-```bash
-msmodelslim quant \
-  --model_path /path/to/wan2_2_t2v_float_weights \
-  --save_path /path/to/wan2_2_t2v_quantized_weights \
-  --device npu \
-  --model_type Wan2_2 \
-  --config_path /lab_practice/wan2_2/wan2_2_w8a8f8_mxfp_t2v.yaml \
-  --trust_remote_code True
-```
-
-After quantization completes, the output directory will contain the quantized model files.
-
-For more examples, refer to the [official examples](https://gitcode.com/Ascend/msit/tree/master/msmodelslim/example).
-
-## Configuration
-
-1. **CLI**: pass `--quantization ascend`.
-
-```bash
-# Offline inference
-python text_to_image.py --model <your-model> --quantization ascend
-
-# Online serving
-vllm serve <your-model> --omni --quantization ascend
-```
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| HunyuanImage-3.0 | - | All layers | None |
-
-Currently, quantized HunyuanImage-3.0 weights have not been uploaded to public model platforms such as Hugging Face. You can use a [HunyuanImage-3.0-adapted msModelSlim version](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8) to generate the quantized weights manually. We will upload the quantized weights as soon as possible.

docs/user_guide/diffusion/quantization/overview.md

@@ -1,153 +0,0 @@
-# Quantization
-
-vLLM-Omni provides a unified quantization framework that supports both diffusion models (DiT, text encoders, VAEs) and multi-stage omni models (e.g., Qwen3-Omni thinker). Quantization reduces memory usage and accelerates inference.
-
-## Supported Methods
-
-### Diffusion Models
-
-| Method | Guide | Description | Tested Models | Min GPU |
-|--------|-------|-------------|---------------|---------|
-| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
-| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
-| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
-| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
-
-### Multi-stage Omni Models (Pre-quantized Checkpoints)
-
-| Method | Description | Tested Models | Min GPU |
-|--------|-------------|---------------|---------|
-| ModelOpt FP8 | Pre-quantized FP8 via NVIDIA ModelOpt | Qwen3-Omni (thinker) | SM 89 (Ada/Hopper) |
-| ModelOpt NVFP4 | Pre-quantized NVFP4 via NVIDIA ModelOpt | Qwen3-Omni (experimental, quality issues) | SM 100 (Blackwell) |
-
-!!! note
-    AWQ, GPTQ, and BitsAndBytes are supported by vLLM's upstream quantization registry but have **not been tested** with vLLM-Omni pipelines. They may work via `build_quant_config()` but are not validated.
-
-### Pre-quantized LLM Checkpoints (Multi-stage Models)
-
-For multi-stage models like Qwen3-Omni, the unified quantization framework auto-detects
-pre-quantized checkpoints via `quantization_config` in the HF config. Quantization is
-automatically scoped to the thinker's `language_model` — audio encoder, vision encoder,
-talker, and code2wav remain in BF16.
-
-| Format | `quant_algo` | Hardware | Example | Status |
-|--------|-------------|----------|---------|--------|
-| ModelOpt FP8 | `FP8` | Ada/Hopper (SM 89+) | `asdazd/Qwen3-Omni-30B-A3B-Instruct_modelopt_FP8` | Tested — 47% memory reduction, comparable throughput |
-| ModelOpt NVFP4 | `NVFP4` | Blackwell (SM 100+) | `shunyang90/Qwen3-Omni-30B-A3B-Instruct-NVFP4` | Experimental — loads and runs but output quality is unacceptable |
-
-**Tested FP8 results (Qwen3-Omni, 1×H200):**
-
-| Config | Model Memory (GiB) | Mem Reduction | Decode (tok/s) |
-|--------|-------------------|---------------|----------------|
-| BF16 baseline | 59.26 | — | 41.6 |
-| FP8 (ModelOpt) | 31.41 | 47% | 39.9 |
-
-FP8 enables the full 3-stage pipeline (thinker + talker + code2wav) on a single 64GB GPU,
-which is impossible with BF16 (thinker alone requires 59.26 GiB).
-
-## Quantization Scope
-
-### Diffusion Models
-
-When `--quantization fp8` is enabled for diffusion models:
-
-| Component | What Gets Quantized | Mechanism |
-|-----------|-------------------|-----------|
-| **DiT (transformer)** | `nn.Linear` layers | vLLM W8A8 FP8 compute (Ada/Hopper) or weight-only (older GPUs) |
-| **Text encoder** | `nn.Linear` layers | FP8 weight storage, BF16 compute |
-| **VAE** | `nn.Conv2d`, `nn.Conv3d` layers | FP8 weight storage, BF16 compute |
-
-### Multi-stage Omni Models
-
-For pre-quantized omni model checkpoints, quantization is scoped per component:
-
-| Component | Quantized? | Notes |
-|-----------|-----------|-------|
-| **Thinker (language_model)** | Yes | FP8 or NVFP4 via ModelOpt |
-| **Audio encoder** | No | Stays BF16 |
-| **Vision encoder** | No | Stays BF16 |
-| **Talker** | No | Stays BF16 |
-| **Code2Wav** | No | Stays BF16 |
-
-!!! note
-    Not all models support all components. See the [FP8 supported models table](fp8.md#supported-models) for per-model details.
-
-## Device Compatibility for FP8
-
-| Device | Example Hardware | FP8 | NVFP4 |
-|--------|-----------------|-----|-------|
-| Blackwell GPU (SM 100+) | RTX 5090, B100, B200 | Yes | Yes (native FP4 HW) |
-| Ada/Hopper GPU (SM 89+) | RTX 4090, H100, H200 | Yes (W8A8 native) | No |
-| Turing/Ampere GPU (SM 75-86) | RTX 3090, A100 | Yes (weight-only Marlin) | No |
-| Ascend NPU | Atlas 800T A2 (910B) | Not yet supported | No |
-
-Kernel selection is automatic on CUDA GPUs.
-
-## Device Compatibility for Int8
-
-| Device Type | Generation | Example | Int8 Mode |
-|-------------|---------------|-------------------|----------|
-| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
-| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
-
-## Python API
-
-The `build_quant_config()` factory accepts multiple input formats:
-
-```python
-from vllm_omni.quantization import build_quant_config
-
-# String
-config = build_quant_config("fp8")
-
-# Dict with parameters
-config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
-
-# AutoRound / INC (auto-detected from checkpoint, or explicit)
-config = build_quant_config("auto-round", bits=4, group_size=128)
-
-# Per-component dict
-config = build_quant_config({
-    "transformer": {"method": "fp8"},
-    "vae": None,
-})
-
-# None / "none"
-config = build_quant_config(None)  # returns None
-```
-
-## Architecture
-
-```
-build_quant_config(spec)
-    ├── str  → _build_single(method) → vLLM registry / _OVERRIDES
-    ├── dict with "method" → _build_single(method, **kwargs)
-    ├── per-component dict → ComponentQuantizationConfig
-    │       routes get_quant_method() by longest-prefix match
-    ├── QuantizationConfig → passthrough
-    └── None / "none" → None
-```
-
-## Migration Guide
-
-### Before (v0.14.x)
-
-```python
-# Old diffusion-specific API (removed)
-from vllm_omni.diffusion.quantization import get_diffusion_quant_config
-config = get_diffusion_quant_config("fp8", activation_scheme="static")
-```
-
-### After (v0.16.0+)
-
-```python
-# Unified API — delegates to vLLM's registry
-from vllm_omni.quantization import build_quant_config
-config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
-
-# Per-component (new)
-config = build_quant_config({
-    "transformer": {"method": "fp8"},
-    "vae": None,
-})
-```

docs/user_guide/diffusion_features.md

@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
 | Method | Description | Best For |
 |--------|-------------|----------|
 | **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
-| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
+| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
 | **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
 
 ### Extensions
@@ -86,9 +86,9 @@ Execution modes control how the diffusion pipeline processes denoise steps.
 
 | Method | Configuration | Description | Best For |
 |--------|--------------|-------------|----------|
-| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
-| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
-| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
+| **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup |
+| **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility |
+| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
 
 ## Supported Models
 
@@ -101,14 +101,14 @@ The following tables show which models support each feature:
 > Notes:
 
 > 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
-> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅  | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
@@ -261,7 +261,7 @@ Measured on NVIDIA H800:
 
 - **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
 - **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
-- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
+- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 
 **Extensions:**
 

docs/user_guide/examples/offline_inference/glm_image.md

@@ -1,154 +1,87 @@
-# GLM-Image Multistage End-to-End Inference
+# GLM-Image Offline Inference
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
-
-
-This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
+declarative config system. The pipeline topology and stage structure are declared in
+`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs live in
+`vllm_omni/deploy/glm_image.yaml`.
 
 ## Architecture
 
-GLM-Image uses a 2-stage pipeline:
-
 ```
-┌─────────────────────────────────────────────────────────────┐
-│                     GLM-Image Pipeline                       │
-├─────────────────────────────────────────────────────────────┤
-│                                                             │
-│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
-│  ┌─────────────────┐            ┌─────────────────────┐     │
-│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
-│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
-│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
-│  │ Generation      │            │  └───────────────┘  │     │
-│  │ (9B AR model)   │            │         │          │     │
-│  └─────────────────┘            │         ▼          │     │
-│         ▲                       │  ┌───────────────┐  │     │
-│         │                       │  │  VAE Decode   │──┼──► Image
-│    Text/Image                   │  └───────────────┘  │     │
-│      Input                      └─────────────────────┘     │
-│                                                             │
-└─────────────────────────────────────────────────────────────┘
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
 ```
 
-## Features
+## Text-to-Image
 
-- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
-- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
-- **Text-to-Image**: Generate images from text descriptions
-- **Image-to-Image**: Edit existing images with text prompts
+```python
+from vllm_omni.entrypoints.omni import Omni
 
-## Usage
-
-### Text-to-Image
-
-```bash
-python end2end.py \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "A beautiful sunset over the ocean with sailing boats" \
-    --height 1024 \
-    --width 1024 \
-    --output output_t2i.png
+if __name__ == "__main__":
+    omni = Omni(model="zai-org/GLM-Image")
+    outputs = omni.generate(
+        "A photorealistic mountain landscape at sunset",
+        sampling_params={
+            "height": 1024,
+            "width": 1024,
+            "num_inference_steps": 50,
+            "guidance_scale": 1.5,
+            "seed": 42,
+        },
+    )
+    outputs[0].request_output.images[0].save("output.png")
 ```
 
-### Image-to-Image (Image Editing)
+## Image-to-Image (Image Editing)
 
-```bash
-python end2end.py \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "Transform this scene into a winter wonderland" \
-    --image input.png \
-    --output output_i2i.png
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="zai-org/GLM-Image")
+    outputs = omni.generate(
+        {
+            "prompt": "Convert this image to watercolor style",
+            "multi_modal_data": {
+                "image": "input.png",
+            },
+        },
+        sampling_params={
+            "height": 1024,
+            "width": 1024,
+            "num_inference_steps": 50,
+            "guidance_scale": 1.5,
+            "seed": 42,
+        },
+    )
+    outputs[0].request_output.images[0].save("output.png")
 ```
 
-### With Custom Parameters
+## Generation Parameters
 
-```bash
-python end2end.py \
-    --model-path /path/to/glm-image \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "A photorealistic cat sitting on a window sill" \
-    --height 1024 \
-    --width 1024 \
-    --num-inference-steps 50 \
-    --guidance-scale 1.5 \
-    --seed 42 \
-    --output output.png
-```
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | None    | Optional random seed                |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
 
-## Shell Scripts
+## VRAM Requirements
 
-### Run Text-to-Image
-
-```bash
-./run_t2i.sh
-```
-
-### Run Image-to-Image
-
-```bash
-./run_i2i.sh --image /path/to/input.png
-```
-
-## Stage Configuration
-
-The stage config (`glm_image.yaml`) defines:
-
-- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
-
-  - Model: `GlmImageForConditionalGeneration`
-  - Output: `token_ids` (prior tokens)
-
-- **Stage 1 (Diffusion)**: Uses diffusion engine
-  - Model: `GlmImagePipeline`
-  - Output: Generated image
-
-See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
-
-## Comparison with Single-Stage
-
-| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
-| ----------- | --------------------------- | ------------------- |
-| AR Model    | transformers native         | vLLM PagedAttention |
-| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
-| Throughput  | Lower                       | Higher              |
-| Flexibility | Single GPU                  | Multi-GPU support   |
-
-## Troubleshooting
-
-### OOM Error
-
-Try reducing memory usage:
-
-```bash
-# In glm_image.yaml, adjust:
-gpu_memory_utilization: 0.5  # Reduce from 0.6
-```
-
-### Slow Initialization
-
-The first run loads model weights. Subsequent runs are faster:
-
-```bash
---stage-init-timeout 900  # Increase timeout for slow storage
-```
-
-## Requirements
-
-- vLLM-Omni with GLM-Image support
-- CUDA-capable GPU (recommended: H100/A100 with 80GB)
-- GLM-Image model weights
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/glm_image/end2end.py"
-    ``````
-??? abstract "run_i2i.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/glm_image/run_i2i.sh"
-    ``````
-??? abstract "run_t2i.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/glm_image/run_t2i.sh"
-    ``````
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |

docs/user_guide/examples/offline_inference/internvla_a1.md

@@ -0,0 +1,205 @@
+# InternVLA-A1
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/internvla_a1>.
+
+
+## Setup
+This example is adapted from: https://github.com/InternRobotics/InternVLA-A1/blob/master/tests/policies/internvla_a1_3b/open_loop_genie1_real.ipynb
+
+This example runs the single-path vLLM offline inference workflow for InternVLA-A1 open-loop action prediction.
+
+Before running the script, export the required local paths:
+
+```bash
+# hf Jia-Zeng/InternVLA-A1-3B-FineTuned-Place_Markpen
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+# hf download InternRobotics/InternData-A1 real_lerobotv30/genie1/Genie1-Place_Markpen.tar.gz --repo-type dataset --local-dir /path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+```
+
+The shell entrypoint also accepts these variables as defaults, so you can keep the command line itself short.
+
+`INTERNVLA_A1_COSMOS_DIR` is expected to contain:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+## Run examples
+
+Get into the example folder:
+
+```bash
+cd examples/offline_inference/internvla_a1
+```
+
+### Run one sample
+
+This runs a single sample through the maintained vLLM registry path.
+
+```bash
+bash run.sh --num-samples 1 --num-episodes 0
+```
+
+### Run open-loop evaluation against GT
+
+This runs one episode, writes evaluation metrics, and saves prediction-vs-GT plots when `matplotlib` is available.
+
+```bash
+bash run.sh --num-episodes 1
+```
+
+Outputs are written under:
+
+```bash
+outputs/internvla_a1/vllm_infer/
+```
+
+Typical files:
+
+- `summary.json`: top-level run summary
+- `registry/log.json`: per-episode GT comparison metrics
+- `registry/plots/*.jpg`: prediction-vs-GT figures
+
+### Eager float32 validation
+
+Use eager attention and float32 when you want the most stable numerical comparison baseline.
+
+```bash
+bash run.sh \
+  --num-episodes 1 \
+  --attn-implementation eager \
+  --dtype float32
+```
+
+### Optional runtime switches
+
+- `--dtype {bfloat16,float32}`: choose inference dtype
+- `--attn-implementation {eager,sdpa}`: switch attention backend
+- `--enable-regional-compile`: enable regional `torch.compile`
+- `--enable-warmup`: run pipeline warmup in initialization
+- `--skip-plots`: skip plot generation even if `matplotlib` is installed
+
+### Collect results and performance logs
+
+Use the helper below when you want to collect one-sample output, forward latency, GT comparison, plots, timing logs, and optional pytest output into a single result directory:
+
+```bash
+bash collect_results.sh
+```
+
+The script writes a timestamped directory under:
+
+```bash
+outputs/internvla_a1/collected_results/
+```
+
+Typical result files:
+
+- `env_summary.txt`: environment and path summary
+- `sample_run.log`: one-sample inference log
+- `forward_benchmark/forward_latency.json`: isolated `pipeline.forward` latency summary
+- `eval_run.log`: open-loop GT evaluation log
+- `eval_outputs/summary.json`: top-level output summary
+- `eval_outputs/registry/log.json`: GT comparison metrics
+- `eval_outputs/registry/plots/*.jpg`: prediction-vs-GT figures
+- `*_time.txt`: timing output
+- `*_gpu.csv`: sampled GPU usage
+
+### Benchmark forward latency
+
+Use the dedicated mode below when you want the isolated `pipeline.forward` latency instead of end-to-end script time:
+
+```bash
+python end2end.py \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir outputs/internvla_a1/forward_benchmark
+```
+
+The output JSON is written to `outputs/internvla_a1/forward_benchmark/forward_latency.json` and contains:
+
+- `cold_start_ms`
+- `warmup_summary`
+- `benchmark_summary`
+- `benchmark_samples_ms`
+
+### Reference results
+
+Reference run collected on `1x NVIDIA H200`, `bfloat16`, `eager`:
+
+- one-sample end-to-end run: `38s`
+- one-episode GT evaluation run: `45s`
+- `average_mse = 1.7173260857816786e-05`
+- `average_mae = 0.0011860118247568607`
+- `average_mse_joint = 7.42028441891307e-06`
+- `average_mae_joint = 0.0010777723509818316`
+- `average_mse_gripper = 8.544408774469048e-05`
+- `average_mae_gripper = 0.0019436875591054559`
+
+The reference plot filename is:
+
+- `eval_outputs/registry/plots/vllm_registry_open_loop_ep0.jpg`
+
+## Measure latency and VRAM
+
+You can collect end-to-end runtime and peak host memory with:
+
+```bash
+/usr/bin/time -v bash run.sh --num-episodes 1
+```
+
+For GPU memory snapshots, use:
+
+```bash
+nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
+```
+
+Run `nvidia-smi` before and during inference to record the peak device memory for your PR report.
+
+## FAQ
+
+If `matplotlib` is missing, evaluation still runs and only plot generation is skipped. To enable plots:
+
+```bash
+pip install matplotlib
+```
+
+## Example materials
+
+- `end2end.py`: main offline inference and GT evaluation entrypoint
+- `run.sh`: shell wrapper with local path env vars
+- `collect_results.sh`: helper to gather result summaries and performance logs into one directory
+- `internvla_a1_common.py`: dataset, evaluation, and plotting helpers
+
+## Embedded source listings
+
+??? abstract "collect_results.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/collect_results.sh"
+    ``````
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/end2end.py"
+    ``````
+??? abstract "internvla_a1_common.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/internvla_a1_common.py"
+    ``````
+??? abstract "run.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/run.sh"
+    ``````
+??? abstract "standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````json
+    --8<-- "examples/offline_inference/internvla_a1/standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````

docs/user_guide/examples/online_serving/glm_image.md

@@ -1,99 +1,96 @@
 # GLM-Image Online Serving
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
+GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
+declarative config system. The pipeline topology and stage structure are declared in
+`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs (GPU placement,
+memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
 
-
-This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
-
-## Run examples (GLM-Image)
-
-**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
-
-### Launch the Server
+## Start Server
 
 ```bash
-# Use default configuration
 vllm serve zai-org/GLM-Image --omni --port 8091
 ```
 
-Or use the convenience script:
+The config system auto-detects the pipeline from the model's `model_index.json` — no
+manual `--stage-configs-path` or `--deploy-config` needed.
+
+By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
+both stages on a single GPU, override per stage:
 
 ```bash
-cd examples/online_serving/glm_image
-bash run_server.sh
+vllm serve zai-org/GLM-Image --omni --port 8091 \
+    --stage-0-devices 0 --stage-1-devices 0
 ```
 
-If you have a custom stage configs file:
-
-```bash
-vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
-```
-
-### Send Requests
-
-Get into the glm_image folder:
-
-```bash
-cd examples/online_serving/glm_image
-```
-
-Send request via Python:
-
-```bash
-python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
-```
-
-The Python client supports the following command-line arguments:
-
-- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
-- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
-- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
-- `--image` (or `-i`): Input image path (for image-to-image editing)
-- `--height`: Image height in pixels (default: 1024)
-- `--width`: Image width in pixels (default: 1024)
-- `--steps`: Number of inference steps (default: 50)
-- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
-- `--seed`: Random seed (default: 42)
-- `--negative`: Negative prompt
-
-## Modality Control
-
-GLM-Image supports **text-to-image** and **image-to-image** modes.
-
-The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
-
-| Mode           | Input        | Output | Description                        |
-| -------------- | ------------ | ------ | ---------------------------------- |
-| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
-| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+## API Calls
 
 ### Text-to-Image
 
 ```bash
-python openai_chat_client.py \
-    --prompt "A photorealistic mountain landscape at sunset" \
-    --height 1024 \
-    --width 1024 \
-    --output landscape.png
-
-# Or use the curl script:
-bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A photorealistic mountain landscape at sunset"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ### Image-to-Image (Image Editing)
 
 ```bash
-python openai_chat_client.py \
-    --prompt "Convert this image to watercolor style" \
-    --image input.png \
-    --output watercolor.png
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {
+        "role": "user",
+        "content": [
+          {"type": "text", "text": "Convert this image to watercolor style"},
+          {"type": "image_url", "image_url": {"url": "data:image/png;base64,$(base64 -w0 input.png)}"}
+        ]
+      }
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
 
-# Or use the curl script:
-bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+### Using the OpenAI Python SDK
+
+```python
+from openai import OpenAI
+import base64
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.chat.completions.create(
+    model="zai-org/GLM-Image",
+    messages=[{"role": "user", "content": "A beautiful sunset over the ocean"}],
+    extra_body={
+        "height": 1024,
+        "width": 1024,
+        "num_inference_steps": 50,
+        "guidance_scale": 1.5,
+        "seed": 42,
+    },
+)
+
+img_url = response.choices[0].message.content[0].image_url.url
+_, b64_data = img_url.split(",", 1)
+with open("output.png", "wb") as f:
+    f.write(base64.b64decode(b64_data))
 ```
 
 For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
@@ -116,7 +113,7 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
 | `width`               | int   | 1024    | Image width in pixels               |
 | `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
 | `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
-| `seed`                | int   | None    | Optional random seed; `/v1/images/*` generates one server-side if omitted |
+| `seed`                | int   | None    | Optional random seed                |
 | `negative_prompt`     | str   | None    | Negative prompt                     |
 
 ## Response Format
@@ -150,13 +147,12 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
 ## Extract Image
 
 ```bash
-# From a saved JSON response
 cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ## Architecture
 
-GLM-Image uses a 2-stage multistage pipeline:
+GLM-Image uses a 2-stage pipeline:
 
 ```
 Stage 0 (AR Model)                Stage 1 (Diffusion)
@@ -181,41 +177,13 @@ Stage 0 (AR Model)                Stage 1 (Diffusion)
 | Stage-1 (DiT+VAE) | **~20 GiB**            |
 | Total             | **~38 GiB + KV Cache** |
 
-## File Description
-
-| File                        | Description                           |
-| --------------------------- | ------------------------------------- |
-| `run_server.sh`             | Server startup script                 |
-| `run_curl_text_to_image.sh` | Text-to-image curl example            |
-| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
-| `openai_chat_client.py`     | Python client (t2i + i2i)             |
-
 ## FAQ
 
-- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the deploy config:
 
 ```yaml
-# In glm_image.yaml, reduce from default 0.6:
+# In vllm_omni/deploy/glm_image.yaml, reduce from default 0.6:
 gpu_memory_utilization: 0.5
 ```
 
 - The first request may be slow due to model warmup. Subsequent requests will be faster.
-
-## Example materials
-
-??? abstract "openai_chat_client.py"
-    ``````py
-    --8<-- "examples/online_serving/glm_image/openai_chat_client.py"
-    ``````
-??? abstract "run_curl_image_edit.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
-    ``````
-??? abstract "run_curl_text_to_image.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
-    ``````
-??? abstract "run_server.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_server.sh"
-    ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -28,11 +28,13 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 ### Launch individual stages (stage-based CLI)
 
 Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
+The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
+`CUDA_VISIBLE_DEVICES`.
 
 **1. Stage 0 (Thinker + API server)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -42,7 +44,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **2. Stage 1 (Talker)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -52,7 +54,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **3. Stage 2 (Code2Wav)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 2 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -75,7 +77,7 @@ that kind of change. Since each command launches one stage, just pass the knob
 directly on that stage command:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -0,0 +1,193 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
+
+This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Run all curl examples
+bash curl_examples.sh
+
+# Or execute directly
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+### Method 2: Using Python Client
+
+```bash
+cd examples/online_serving/stable_audio
+
+# Simple generation
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+
+# With custom duration
+python stable_audio_client.py \
+    --text "A dog barking" \
+    --audio_length 5.0
+
+# With all parameters
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+The Python client supports the following command-line arguments:
+
+- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
+- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
+- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
+- `--audio_start`: Audio start time in seconds (default: `0.0`)
+- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
+- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
+- `--num_inference_steps`: Number of inference steps (default: `100`)
+- `--seed`: Random seed for reproducibility (default: `None`)
+- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
+- `--output`: Output file path (default: `stable_audio_output.wav`)
+
+### Method 3: Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "The sound of ocean waves crashing on a beach",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+    },
+    timeout=300.0,
+)
+
+with open("ocean.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Request Format
+
+### Simple Generation
+
+```json
+{
+    "input": "The sound of ocean waves"
+}
+```
+
+### Generation with Parameters
+
+```json
+{
+    "input": "A piano playing a gentle melody",
+    "audio_length": 10.0,
+    "negative_prompt": "Low quality, distorted, noisy",
+    "guidance_scale": 8.0,
+    "num_inference_steps": 150,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Generation Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+## Tuning Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality (quick previews)
+   - 100 steps: Good balance (general purpose)
+   - 150+ steps: High quality, slower (final / critical audio)
+3. **Guidance Scale**:
+   - Lower (3 - 5): More creative / varied output
+   - Default (7): Good balance
+   - Higher (10+): Strict adherence to the prompt
+4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
+5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
+
+## File Description
+
+| File | Description |
+|------|-------------|
+| `curl_examples.sh` | Curl examples covering common use cases |
+| `stable_audio_client.py` | Python client with full CLI argument support |
+
+## Troubleshooting
+
+1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
+2. **Connection refused**: Make sure the server is running on the correct port.
+3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
+4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
+5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
+
+## Example materials
+
+??? abstract "stable_audio_client.py"
+    ``````py
+    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
+    ``````
+??? abstract "curl_examples.sh"
+    ``````sh
+    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
+    ``````

docs/user_guide/quantization/autoround.md

@@ -0,0 +1,127 @@
+# AutoRound Quantization
+
+## Overview
+
+[AutoRound](https://github.com/intel/auto-round) produces pre-quantized
+checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
+checkpoint's `config.json` and auto-detects
+`quantization_config.quant_method = "auto-round"`.
+
+AutoRound is static quantization: no `--quantization` flag is needed at
+inference time when the checkpoint already contains the quantization config.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ✅ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. AutoRound is Intel-supported.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
+| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Language-model stage | W4A16 | AutoRound |
+| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound |
+| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
+
+AutoRound support is checkpoint-driven. A model is supported when its
+checkpoint uses a compatible INC/AutoRound config and the target stage maps to
+vLLM-Omni's runtime module names.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| 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
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=28),
+)
+outputs[0].save_images("output.png")
+```
+
+CLI:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
+  --prompt "A cat sitting on a windowsill" \
+  --num-inference-steps 28 \
+  --output outputs/flux_w4a16.png
+```
+
+## Parameters
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `quant_method` | str | Must be `"auto-round"` |
+| `bits` | int | Quantized weight bit width, usually `4` |
+| `group_size` | int | Quantization group size |
+| `packing_format` | str | AutoRound packing format, for example `auto_round:auto_gptq` |
+| `block_name_to_quantize` | str | Checkpoint block names that should map to runtime module names |
+
+The checkpoint should contain a config like:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "auto-round",
+    "bits": 4,
+    "group_size": 128,
+    "sym": true,
+    "packing_format": "auto_round:auto_gptq",
+    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
+  }
+}
+```
+
+## Validation and Notes
+
+At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
+to runtime module names, and selects the matching vLLM compute backend.
+
+Example checkpoint creation:
+
+```bash
+auto-round \
+  --model black-forest-labs/FLUX.1-dev \
+  --scheme W4A16 \
+  --batch_size 1 \
+  --disable_opt_rtn \
+  --dataset coco2014 \
+  --iters 0
+```
+
+Use the generated output directory directly as the `model` argument. See the
+[AutoRound documentation](https://github.com/intel/auto-round) for all
+available schemes and options.

docs/user_guide/quantization/fp8.md

@@ -0,0 +1,121 @@
+# FP8 Quantization
+
+## Overview
+
+FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
+a checkpoint whose target stage already declares an FP8 quantization config.
+Online activation scaling is the default and does not require calibration.
+Static activation scaling is supported when calibrated scale information is
+available.
+
+Some architectures can quantize all linear layers. Others have
+quality-sensitive layers that should stay in BF16 through `ignored_layers`.
+Image-stream MLPs (`img_mlp`) are a common sensitive target because denoising
+latent ranges shift across timesteps and small per-layer errors can compound
+in deep DiT blocks.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` |
+|-------|-----------|:-------:|:------:|----------------|------------------|
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None |
+| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None |
+| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None |
+| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Format | Status |
+|-------|-------|--------|--------|
+| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
+| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
+
+Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
+model-specific guide says otherwise.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="<your-model>", quantization="fp8")
+
+omni_with_skips = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "fp8",
+        "ignored_layers": ["img_mlp"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model <your-model> --quantization fp8
+python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
+vllm serve <your-model> --omni --quantization fp8
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"fp8"`) |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available |
+| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
+
+The available `ignored_layers` names depend on the model architecture, for
+example `to_qkv`, `to_out`, `img_mlp`, or `txt_mlp`.
+
+## Validation and Notes
+
+FP8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="fp8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Compare generated outputs with a BF16 baseline before adding a new model to the
+supported table. GLM-Image and Helios are not listed as FP8-supported diffusion
+models until they have method-specific validation.

docs/user_guide/quantization/gguf.md

@@ -0,0 +1,107 @@
+# GGUF Quantization
+
+## Overview
+
+GGUF support loads pre-quantized diffusion transformer weights while keeping
+the rest of the pipeline on the base Hugging Face checkpoint. Use the base
+model for tokenizer, text encoder, scheduler, and VAE, then pass the GGUF file
+for the transformer.
+
+GGUF is static quantization: the quantized weights are produced before serving.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF base model | GGUF input | Scope | Adapter |
+|-------|---------------|------------|-------|---------|
+| Qwen-Image family | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512`, edit and layered Qwen-Image pipelines | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `QwenImageGGUFAdapter` |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Transformer only | No validated adapter listed |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `ZImageGGUFAdapter` |
+| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `Flux2KleinGGUFAdapter` |
+
+Generic FLUX.1 GGUF checkpoints are not listed here; the implemented adapter is
+for the FLUX.2-klein path.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | GGUF is not documented for omni/TTS AR stages |
+| Qwen3-TTS | TTS language-model stage | Not validated | GGUF is not documented for TTS stages |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
+| GLM-Image | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
+
+## Configuration
+
+Offline:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
+  --quantization gguf \
+  --prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
+  --height 1024 \
+  --width 1024 \
+  --seed 42 \
+  --num_inference_steps 20 \
+  --output outputs/qwen_image_q4km.png
+```
+
+Online:
+
+```bash
+vllm serve Qwen/Qwen-Image \
+  --omni \
+  --port 8000 \
+  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"gguf"`) |
+| `gguf_model` | str | - | Local GGUF file, explicit Hugging Face file, or `repo:quant_type` selector |
+
+`gguf_model` accepts:
+
+| Form | Example |
+|------|---------|
+| Local file | `/models/z-image-Q4_K_M.gguf` |
+| Explicit HF file | `QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf` |
+| HF repo plus quant type | `owner/repo:Q4_K_M` |
+
+## Validation and Notes
+
+1. `OmniDiffusionConfig` receives `{"method": "gguf", "gguf_model": ...}`.
+2. `DiffusersPipelineLoader` resolves the GGUF file.
+3. A model-specific adapter remaps GGUF tensor names to vLLM-Omni transformer
+   names.
+4. Only transformer weights are loaded from GGUF. Missing non-transformer
+   weights are loaded from the base model repository.
+5. vLLM's GGUF linear method performs dequantization and GEMM at runtime.
+
+Unsupported models fail fast with a clear "No GGUF adapter matched" error
+instead of falling back to a generic mapper. Many GGUF repositories do not
+include `model_index.json`; always pass the normal base model through `--model`.

docs/user_guide/quantization/int8.md

@@ -0,0 +1,107 @@
+# Int8 Quantization
+
+## Overview
+
+Int8 quantization supports W8A8 diffusion transformer inference on CUDA and
+Ascend NPU. It can quantize BF16/FP16 weights at load time, or load serialized
+Int8 checkpoints that already contain quantized weights and scales.
+
+Only online activation scaling is currently supported.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | CUDA | Ascend NPU | Mode | Recommendation |
+|-------|-----------|:----:|:----------:|------|----------------|
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Online W8A8 | All layers |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Online W8A8 | Validate before enabling in docs |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | Online W8A8 | All layers |
+
+Other diffusion models may work if their transformer uses supported linear
+layers, but they are not validated in this guide.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
+| Qwen3-TTS | TTS language-model stage | Not validated | No Int8 TTS stage support is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="<your-model>", quantization="int8")
+
+omni_with_skips = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "int8",
+        "ignored_layers": ["<layer-name>"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model <your-model> --quantization int8
+python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
+vllm serve <your-model> --omni --quantization int8
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"int8"`) |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling; static is not supported |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+| `is_checkpoint_int8_serialized` | bool | `False` | Set by checkpoint config when loading serialized Int8 weights |
+
+## Validation and Notes
+
+Int8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="int8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Only add a new model to the supported table after comparing the Int8 output
+against a BF16 baseline and documenting any required `ignored_layers`.

docs/user_guide/quantization/msmodelslim.md

@@ -0,0 +1,95 @@
+# msModelSlim Quantization
+
+## Overview
+
+[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
+toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
+checkpoints run through the Ascend/NPU path with `--quantization ascend`.
+
+msModelSlim is static quantization: quantized weights are generated offline
+before vLLM-Omni inference starts.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ❌ |
+| AMD ROCm | ❌ |
+| Intel XPU | ❌ |
+| Ascend NPU | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | Base model | Scope | Hardware | Notes |
+|-------|------------|-------|----------|-------|
+| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Ascend A2/A3 NPU | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation |
+
+Public Hugging Face quantized weights are not available yet. Use the
+[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
+to generate the checkpoint manually.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker or language-model stage | Not validated | No msModelSlim omni checkpoint path is documented |
+| Qwen3-TTS | TTS language-model stage | Not validated | No msModelSlim TTS checkpoint path is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
+| GLM-Image | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
+
+## Configuration
+
+Offline inference:
+
+```bash
+python text_to_image.py --model <quantized-model-path> --quantization ascend
+```
+
+Online serving:
+
+```bash
+vllm serve <quantized-model-path> --omni --quantization ascend
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--quantization` | str | - | Use `ascend` for msModelSlim-produced checkpoints |
+| `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling |
+
+Example msModelSlim command for a Wan2.2 W8A8 checkpoint:
+
+```bash
+msmodelslim quant \
+  --model_path /path/to/wan2_2_t2v_float_weights \
+  --save_path /path/to/wan2_2_t2v_quantized_weights \
+  --device npu \
+  --model_type Wan2_2 \
+  --config_path /path/to/wan2_2_w8a8f8_mxfp_t2v.yaml \
+  --trust_remote_code True
+```
+
+For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
+
+## Validation and Notes
+
+1. Run with the Ascend/NPU installation and environment.
+2. The `ascend` quantization method expects weights produced by the Ascend
+   tooling; it is not a load-time CUDA quantizer.
+3. Keep the quantized checkpoint aligned with the same model architecture and
+   stage config used for BF16 inference.

docs/user_guide/quantization/online.md

@@ -0,0 +1,93 @@
+# Online Quantization
+
+## Overview
+
+Online quantization means vLLM-Omni computes quantized weights and scales while
+loading the model. Use it when you want memory savings without preparing a
+separate quantized checkpoint.
+
+This mode is different from pre-quantized checkpoint formats such as GGUF,
+AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
+prepared before serving and are documented in their method-specific guides.
+
+## Hardware Support
+
+| Device | FP8 W8A8 | Int8 W8A8 |
+|--------|----------|-----------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
+| AMD ROCm | ⭕ | ⭕ |
+| Intel XPU | ⭕ | ⭕ |
+| Ascend NPU | ❌ | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Method | Guide | Example models | Status |
+|--------|-------|----------------|--------|
+| 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 |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+Online quantization is not currently validated for the omni/TTS stages. For
+Qwen3-Omni and related models, prefer checkpoint-declared ModelOpt or
+AutoRound paths when available.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+Online quantization must be routed to the intended stage. BAGEL and GLM-Image
+need model-specific validation before they are listed as supported targets.
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+
+omni_fp8 = Omni(model="<your-model>", quantization="fp8")
+omni_int8 = Omni(model="<your-model>", quantization="int8")
+```
+
+CLI:
+
+```bash
+vllm serve <your-model> --omni --quantization fp8
+vllm serve <your-model> --omni --quantization int8
+```
+
+Per-component routing:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+```
+
+## Parameters
+
+| Parameter | Methods | Description |
+|-----------|---------|-------------|
+| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
+| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
+| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
+| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
+
+## Validation and Notes
+
+1. Compare the online-quantized output against a BF16 baseline with the same
+   seed and generation parameters.
+2. Use `ignored_layers` for quality-sensitive MLPs or output projections.
+3. Document any required skipped layers in the method page before marking a new
+   model as supported.
+4. If a model already ships quantized weights, use the matching pre-quantized
+   method guide instead of online quantization.

docs/user_guide/quantization/overview.md

@@ -0,0 +1,134 @@
+# Quantization
+
+vLLM-Omni exposes quantization through the unified `quantization_config`
+path. The same configuration entrypoint is used across diffusion-only models,
+multi-stage omni/TTS models, and multi-stage diffusion models, but each model
+type has a different quantization scope.
+
+## Quantization Modes
+
+| Mode | Guide | Description | Methods |
+|------|-------|-------------|---------|
+| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
+| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |
+
+## Hardware Support
+
+| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
+|--------|----------|-----------|------|-----------|-------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+These models run a diffusion transformer as the primary inference module. The
+default quantization target is the transformer; tokenizer, scheduler, text
+encoder, and VAE stay on the base checkpoint unless a method guide says
+otherwise.
+
+| Method | Guide | Mode | Example models | Status |
+|--------|-------|------|----------------|--------|
+| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
+| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
+| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
+| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+These models combine an AR language model with audio, vision, talker, or TTS
+stages. Quantization is scoped to the AR language-model stage when the
+checkpoint contains a supported `quantization_config`; the non-AR stages stay
+in BF16 unless the model guide explicitly adds support.
+
+| Method | Guide | Scope | Example models | Status |
+|--------|-------|-------|----------------|--------|
+| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
+| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+These models split generation across multiple stages. Quantization must be
+attached to the intended stage rather than applied globally.
+
+| Method | Guide | Scope | Example models | Status |
+|--------|-------|-------|----------------|--------|
+| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
+| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
+| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
+| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
+
+!!! note
+    "Online quantization" means vLLM-Omni computes the quantization data while
+    loading the model. "Pre-quantized" means the checkpoint or external
+    quantizer provides the required quantized weights and scales.
+
+## Quantization Scope
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+The default target is the diffusion transformer. Component routing is available
+through `build_quant_config()`:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+```
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
+| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
+| VAE | No | Keep BF16; storage-only paths are method-specific |
+| Scheduler/tokenizer | No | Loaded from the base model repository |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Thinker or AR language model | Yes, when checkpoint config is supported | ModelOpt FP8/NVFP4 or AutoRound checkpoint config |
+| Audio encoder | No | BF16 |
+| Vision encoder | No | BF16 |
+| Talker or TTS stage | No | BF16 unless model-specific support is documented |
+| Code2Wav | No | BF16 |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Selected diffusion or transformer stage | Method-specific | Must be routed to the intended stage |
+| Other generation stages | No | Keep BF16 unless separately validated |
+| VAE, tokenizer, scheduler | No | Loaded from the base checkpoint |
+
+## Python API
+
+`build_quant_config()` accepts strings, dictionaries, per-component
+dictionaries, existing `QuantizationConfig` objects, or `None`.
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+build_quant_config("fp8")
+build_quant_config({"method": "fp8", "activation_scheme": "static"})
+build_quant_config("auto-round", bits=4, group_size=128)
+build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
+build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
+build_quant_config(None)
+```

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -118,9 +118,6 @@ def parse_args() -> argparse.Namespace:
         help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/helios/end2end.py

@@ -196,9 +196,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
     parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/image_to_image/image_edit.py

@@ -378,9 +378,6 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -211,9 +211,6 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/internvla_a1/README.md

@@ -0,0 +1,30 @@
+# InternVLA-A1
+
+Full usage and result-reporting guidance lives in [docs/user_guide/examples/offline_inference/internvla_a1.md](../../../docs/user_guide/examples/offline_inference/internvla_a1.md).
+
+Quick start:
+
+```bash
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+
+bash run.sh --num-samples 1 --num-episodes 0
+bash run.sh --num-episodes 1
+bash collect_results.sh
+```
+
+Expected files under `INTERNVLA_A1_COSMOS_DIR`:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+Key entrypoints:
+
+- `run.sh`: wrapper for offline inference and GT evaluation
+- `collect_results.sh`: collect sample output, latency, metrics, plots, and logs
+- `end2end.py`: underlying Python entrypoint

examples/offline_inference/internvla_a1/collect_results.sh

@@ -0,0 +1,220 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+RESULT_ROOT="${INTERNVLA_A1_RESULT_DIR:-$REPO_ROOT/outputs/internvla_a1/collected_results/$TIMESTAMP}"
+
+: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+
+EVAL_OUTPUT_DIR="$RESULT_ROOT/eval_outputs"
+GPU_MONITOR_INTERVAL="${GPU_MONITOR_INTERVAL:-1}"
+
+mkdir -p "$RESULT_ROOT"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+write_env_summary() {
+  cat >"$RESULT_ROOT/env_summary.txt" <<EOF
+timestamp=$TIMESTAMP
+repo_root=$REPO_ROOT
+result_root=$RESULT_ROOT
+python=$(command -v python || true)
+python_version=$(python --version 2>&1 || true)
+model_dir=$INTERNVLA_A1_MODEL_DIR
+dataset_dir=$INTERNVLA_A1_DATASET_DIR
+processor_dir=$INTERNVLA_A1_PROCESSOR_DIR
+cosmos_dir=$INTERNVLA_A1_COSMOS_DIR
+pwd=$(pwd)
+EOF
+}
+
+capture_gpu_snapshot() {
+  local output_file="$1"
+  if command -v nvidia-smi >/dev/null 2>&1; then
+    nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv >"$output_file" 2>&1 || true
+  else
+    echo "nvidia-smi not found" >"$output_file"
+  fi
+}
+
+start_gpu_monitor() {
+  local output_file="$1"
+  if ! command -v nvidia-smi >/dev/null 2>&1; then
+    echo "nvidia-smi not found" >"$output_file"
+    return 1
+  fi
+
+  (
+    echo "timestamp,index,name,memory.total [MiB],memory.used [MiB],utilization.gpu [%]"
+    while true; do
+      local now
+      now="$(date '+%Y-%m-%d %H:%M:%S')"
+      nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv,noheader,nounits \
+        | awk -F',' -v ts="$now" '{gsub(/^[ \t]+|[ \t]+$/, "", $0); print ts "," $0}'
+      sleep "$GPU_MONITOR_INTERVAL"
+    done
+  ) >"$output_file" 2>/dev/null &
+  echo $!
+}
+
+stop_gpu_monitor() {
+  local monitor_pid="${1:-}"
+  if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then
+    kill "$monitor_pid" 2>/dev/null || true
+    wait "$monitor_pid" 2>/dev/null || true
+  fi
+}
+
+copy_eval_outputs() {
+  local source_dir="$1"
+  local target_dir="$2"
+  local source_real=""
+  local target_real=""
+  source_real="$(realpath "$source_dir" 2>/dev/null || echo "$source_dir")"
+  target_real="$(realpath "$target_dir" 2>/dev/null || echo "$target_dir")"
+  if [[ "$source_real" == "$target_real" ]]; then
+    return 0
+  fi
+  mkdir -p "$target_dir"
+  if [[ -d "$source_dir" ]]; then
+    cp -r "$source_dir"/. "$target_dir"/
+  fi
+}
+
+run_with_artifacts() {
+  local name="$1"
+  shift
+
+  local log_file="$RESULT_ROOT/${name}.log"
+  local time_file="$RESULT_ROOT/${name}_time.txt"
+  local gpu_file="$RESULT_ROOT/${name}_gpu.csv"
+  local status_file="$RESULT_ROOT/${name}_status.txt"
+  local monitor_pid=""
+
+  monitor_pid="$(start_gpu_monitor "$gpu_file" || true)"
+  set +e
+  local exit_code=0
+  local start_ts=0
+  local end_ts=0
+  start_ts="$(date +%s)"
+  if [[ -x /usr/bin/time ]]; then
+    /usr/bin/time -v -o "$time_file" "$@" >"$log_file" 2>&1
+    exit_code=$?
+    if [[ ! -s "$time_file" ]]; then
+      end_ts="$(date +%s)"
+      {
+        echo "timing_mode=usr_bin_time_empty_fallback"
+        echo "start_ts=$start_ts"
+        echo "end_ts=$end_ts"
+        echo "elapsed_seconds=$((end_ts - start_ts))"
+      } >"$time_file"
+    fi
+  else
+    "$@" >"$log_file" 2>&1
+    exit_code=$?
+    end_ts="$(date +%s)"
+    {
+      echo "timing_mode=shell_date"
+      echo "start_ts=$start_ts"
+      echo "end_ts=$end_ts"
+      echo "elapsed_seconds=$((end_ts - start_ts))"
+    } >"$time_file"
+  fi
+  set -e
+  stop_gpu_monitor "$monitor_pid"
+
+  echo "exit_code=$exit_code" >"$status_file"
+  if [[ $exit_code -ne 0 ]]; then
+    echo "[error] command failed for $name, see $log_file"
+    return $exit_code
+  fi
+}
+
+write_manifest() {
+  cat >"$RESULT_ROOT/README.txt" <<EOF
+InternVLA-A1 collected results
+
+Key files:
+- env_summary.txt: environment and path summary
+- sample_run.log / sample_run_time.txt: one-sample functional run
+- forward_benchmark.log / forward_benchmark_time.txt: pure pipeline.forward latency benchmark
+- eval_run.log / eval_run_time.txt: GT evaluation run
+- pytest_e2e.log / pytest_e2e_time.txt: offline e2e pytest result
+- gpu_info_before.csv / gpu_info_after.csv: point-in-time GPU snapshots
+- *_gpu.csv: sampled GPU usage during each run
+- eval_outputs/: copied output directory from the GT evaluation run
+
+Important outputs:
+- forward_benchmark/forward_latency.json
+- eval_outputs/summary.json
+- eval_outputs/registry/log.json
+- eval_outputs/registry/plots/
+EOF
+}
+
+write_skip_artifact() {
+  local name="$1"
+  local reason="$2"
+  echo "$reason" >"$RESULT_ROOT/${name}.log"
+  echo "timing_mode=skipped" >"$RESULT_ROOT/${name}_time.txt"
+  echo "exit_code=0" >"$RESULT_ROOT/${name}_status.txt"
+  echo "skipped_reason=$reason" >>"$RESULT_ROOT/${name}_status.txt"
+  if [[ ! -f "$RESULT_ROOT/${name}_gpu.csv" ]]; then
+    echo "skipped,$reason" >"$RESULT_ROOT/${name}_gpu.csv"
+  fi
+}
+
+write_env_summary
+write_manifest
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_before.csv"
+
+run_with_artifacts \
+  "sample_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$RESULT_ROOT/sample_outputs" \
+  --num-samples 1 \
+  --num-episodes 0
+
+run_with_artifacts \
+  "forward_benchmark" \
+  python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir "$RESULT_ROOT/forward_benchmark"
+
+run_with_artifacts \
+  "eval_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$EVAL_OUTPUT_DIR" \
+  --num-episodes 1
+
+copy_eval_outputs "$EVAL_OUTPUT_DIR" "$RESULT_ROOT/eval_outputs"
+
+if python - <<'PY' >/dev/null 2>&1
+import importlib.util
+raise SystemExit(0 if importlib.util.find_spec("pytest") else 1)
+PY
+then
+  run_with_artifacts \
+    "pytest_e2e" \
+    python -m pytest -sv tests/e2e/offline_inference/test_internvla_a1.py -m advanced_model
+else
+  write_skip_artifact "pytest_e2e" "pytest is not installed in the current python environment"
+fi
+
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_after.csv"
+
+echo "Results written to: $RESULT_ROOT"

examples/offline_inference/internvla_a1/end2end.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parent
+_REPO_ROOT = _WORKSPACE_ROOT.parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    # Keep the shared helper import clean by doing script-only path bootstrap here.
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from internvla_a1_common import (  # noqa: E402
+    A2DOpenLoopDataset,
+    collate_open_loop_samples,
+    make_shared_noise,
+    run_open_loop_evaluation,
+    select_indices,
+    tensor_dtype,
+    tensor_sha256,
+)
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig  # noqa: E402
+from vllm_omni.diffusion.models.internvla_a1 import (  # noqa: E402
+    InternVLAA1Config,
+    InternVLAA1TrainMetadata,
+)
+from vllm_omni.diffusion.models.internvla_a1.config import OBS_STATE  # noqa: E402
+from vllm_omni.diffusion.registry import initialize_model  # noqa: E402
+from vllm_omni.diffusion.request import OmniDiffusionRequest  # noqa: E402
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams  # noqa: E402
+
+
+def _required_path_arg(env_name: str, cli_value: str | None) -> str:
+    value = cli_value or os.getenv(env_name)
+    if not value:
+        raise ValueError(f"Missing required path: set --{env_name.lower().replace('_', '-')} or {env_name}.")
+    return value
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run vLLM InternVLA-A1 inference on a few samples.")
+    parser.add_argument("--model-dir")
+    parser.add_argument("--dataset-dir")
+    parser.add_argument("--num-samples", type=int, default=1)
+    parser.add_argument("--num-episodes", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--compile-model", action="store_true")
+    parser.add_argument("--attn-implementation", default="eager")
+    parser.add_argument("--enable-regional-compile", action="store_true")
+    parser.add_argument("--enable-warmup", action="store_true")
+    parser.add_argument("--strict-load", action="store_true")
+    parser.add_argument("--output-dir", default="outputs/internvla_a1/vllm_infer")
+    parser.add_argument("--skip-plots", action="store_true")
+    parser.add_argument("--benchmark-forward", action="store_true")
+    parser.add_argument("--warmup-iters", type=int, default=3)
+    parser.add_argument("--benchmark-iters", type=int, default=10)
+    return parser.parse_args()
+
+
+def build_od_config(args: argparse.Namespace, processor_model_name: str) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model=str(Path(args.model_dir).resolve()),
+        model_class_name="InternVLAA1Pipeline",
+        dtype=tensor_dtype(args.dtype),
+        custom_pipeline_args={
+            "device": args.device,
+            "dtype": args.dtype,
+            "compile_model": args.compile_model,
+            "attn_implementation": args.attn_implementation,
+            "enable_regional_compile": args.enable_regional_compile,
+            "enable_warmup": args.enable_warmup,
+            "strict_load": args.strict_load,
+            "processor_model_name": processor_model_name,
+        },
+    )
+
+
+def build_dataset(args: argparse.Namespace) -> tuple[A2DOpenLoopDataset, InternVLAA1Config, InternVLAA1TrainMetadata]:
+    model_dir = Path(args.model_dir)
+    config = InternVLAA1Config.from_pretrained(model_dir)
+    config.device = args.device
+    config.dtype = args.dtype
+    config.compile_model = args.compile_model
+    config.attn_implementation = args.attn_implementation
+    config.enable_regional_compile = args.enable_regional_compile
+
+    train_meta = InternVLAA1TrainMetadata.from_pretrained(model_dir)
+    with open(model_dir / "stats.json", encoding="utf-8") as f:
+        train_stats = json.load(f)["a2d"]
+
+    dataset = A2DOpenLoopDataset(
+        args.dataset_dir,
+        config=config,
+        train_stats=train_stats,
+    )
+    return dataset, config, train_meta
+
+
+def run_one_path(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    indices: list[int],
+) -> list[dict[str, object]]:
+    results: list[dict[str, object]] = []
+    for index in indices:
+        sample = dataset.get_sample(index)
+        batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+        noise = make_shared_noise(
+            args.seed,
+            index,
+            (
+                batch_inputs[OBS_STATE].shape[0],
+                config.chunk_size,
+                config.max_action_dim,
+            ),
+            args.device,
+        )
+        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+        results.append(
+            {
+                "path": "registry",
+                "index": index,
+                "episode_index": sample.episode_index,
+                "task": sample.task,
+                "seed": args.seed,
+                "shape": list(pred.shape),
+                "mean": float(pred.mean().item()),
+                "std": float(pred.std().item()),
+                "action_sha256": tensor_sha256(pred),
+                "first_action_prefix": pred[0, 0, :8].tolist(),
+            }
+        )
+    return results
+
+
+def run_pipeline_forward(
+    pipeline,
+    batch_inputs: dict[str, torch.Tensor],
+    noise: torch.Tensor,
+) -> torch.Tensor:
+    output = pipeline.forward(
+        OmniDiffusionRequest(
+            prompts=[""],
+            sampling_params=OmniDiffusionSamplingParams(
+                extra_args={
+                    "batch_inputs": batch_inputs,
+                    "noise": noise,
+                    "decode_image": False,
+                }
+            ),
+        )
+    )
+    if output.error:
+        raise RuntimeError(output.error)
+    if output.output is None:
+        raise RuntimeError("InternVLAA1Pipeline.forward returned no output tensor.")
+    return output.output
+
+
+def _synchronize(device: str) -> None:
+    if device.startswith("cuda") and torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _latency_summary(values_ms: list[float]) -> dict[str, float]:
+    sorted_values = sorted(values_ms)
+    p50_index = min(len(sorted_values) - 1, round(0.50 * (len(sorted_values) - 1)))
+    p90_index = min(len(sorted_values) - 1, round(0.90 * (len(sorted_values) - 1)))
+    return {
+        "mean_ms": float(statistics.mean(sorted_values)),
+        "stdev_ms": float(statistics.pstdev(sorted_values)) if len(sorted_values) > 1 else 0.0,
+        "min_ms": float(sorted_values[0]),
+        "max_ms": float(sorted_values[-1]),
+        "p50_ms": float(sorted_values[p50_index]),
+        "p90_ms": float(sorted_values[p90_index]),
+    }
+
+
+def benchmark_forward(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    index: int,
+    output_dir: Path,
+) -> dict[str, object]:
+    sample = dataset.get_sample(index)
+    batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+    noise = make_shared_noise(
+        args.seed,
+        index,
+        (
+            batch_inputs[OBS_STATE].shape[0],
+            config.chunk_size,
+            config.max_action_dim,
+        ),
+        args.device,
+    )
+
+    _synchronize(args.device)
+    cold_start_begin = time.perf_counter()
+    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 _ in range(args.warmup_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        warmup_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    benchmark_ms: list[float] = []
+    for _ in range(args.benchmark_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+    summary = {
+        "mode": "forward_latency",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "sample_index": index,
+        "episode_index": sample.episode_index,
+        "task": sample.task,
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "warmup_iters": args.warmup_iters,
+        "benchmark_iters": args.benchmark_iters,
+        "output_shape": list(pred.shape),
+        "cold_start_ms": cold_start_ms,
+        "warmup_summary": _latency_summary(warmup_ms) if warmup_ms else {},
+        "benchmark_summary": _latency_summary(benchmark_ms) if benchmark_ms else {},
+        "benchmark_samples_ms": benchmark_ms,
+    }
+    with open(output_dir / "forward_latency.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+    return summary
+
+
+def main() -> None:
+    args = parse_args()
+    args.model_dir = _required_path_arg("INTERNVLA_A1_MODEL_DIR", args.model_dir)
+    args.dataset_dir = _required_path_arg("INTERNVLA_A1_DATASET_DIR", args.dataset_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset, config, train_meta = build_dataset(args)
+    indices = select_indices(dataset, args.num_samples)
+    od_config = build_od_config(args, train_meta.processor_model_name)
+
+    eval_summaries: dict[str, object] = {}
+    pipeline = initialize_model(od_config)
+    if args.benchmark_forward:
+        benchmark_forward(
+            pipeline=pipeline,
+            dataset=dataset,
+            config=config,
+            args=args,
+            index=indices[0],
+            output_dir=output_dir,
+        )
+        return
+    results = run_one_path(
+        pipeline=pipeline,
+        dataset=dataset,
+        config=config,
+        args=args,
+        indices=indices,
+    )
+    if args.num_episodes > 0:
+        eval_summaries["registry"] = run_open_loop_evaluation(
+            mode="vllm_registry",
+            policy=pipeline,
+            config=config,
+            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),
+            num_episodes=args.num_episodes,
+            seed=args.seed,
+            device=args.device,
+            dtype=tensor_dtype(args.dtype),
+            output_dir=output_dir / "registry",
+            skip_plots=args.skip_plots,
+        )
+
+    summary = {
+        "mode": "registry",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "seed": args.seed,
+        "indices": indices,
+        "results": results,
+        "output_dir": str(output_dir.resolve()),
+        "eval_summaries": eval_summaries,
+    }
+    with open(output_dir / "summary.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/internvla_a1/internvla_a1_common.py

@@ -0,0 +1,471 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import torch.nn.functional as F
+import torchvision
+
+from vllm_omni.diffusion.models.internvla_a1.config import (  # noqa: E402
+    OBS_IMAGES,
+    OBS_STATE,
+    OBS_TASK,
+    InternVLAA1Config,
+)
+
+
+def tensor_dtype(name: str) -> torch.dtype:
+    return torch.bfloat16 if name == "bfloat16" else torch.float32
+
+
+def select_indices(dataset, num_samples: int) -> list[int]:
+    indices: list[int] = []
+    for _, episode_indices in dataset.episode_start_indices(max_episodes=num_samples):
+        if episode_indices:
+            indices.append(episode_indices[0])
+    return indices[:num_samples]
+
+
+def make_shared_noise(seed: int, sample_index: int, shape: tuple[int, ...], device: str) -> torch.Tensor:
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed + sample_index)
+    noise = torch.randn(shape, generator=generator, dtype=torch.float32)
+    return noise.to(device=device)
+
+
+def tensor_sha256(tensor: torch.Tensor) -> str:
+    array = tensor.detach().contiguous().cpu().numpy()
+    return hashlib.sha256(array.tobytes()).hexdigest()
+
+
+def _clamp_index(index: int, start: int, end: int) -> int:
+    return max(start, min(end - 1, index))
+
+
+def _load_parquet_rows(path: Path) -> list[dict[str, Any]]:
+    return pq.read_table(path).to_pylist()
+
+
+def _load_json(path: Path) -> dict[str, Any]:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _stack_stats(stats: dict[str, Any], keys: list[str]) -> dict[str, torch.Tensor]:
+    result = {}
+    for stat_name in ("mean", "std"):
+        values = []
+        for key in keys:
+            values.extend(stats[key][stat_name])
+        result[stat_name] = torch.tensor(values, dtype=torch.float32)
+    return result
+
+
+def normalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    denom = torch.where(stats["std"] == 0, torch.ones_like(stats["std"]), stats["std"])
+    return (values - stats["mean"]) / denom
+
+
+def unnormalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    return values * stats["std"] + stats["mean"]
+
+
+class TorchvisionVideoReaderCache:
+    def __init__(self, backend: str = "pyav") -> None:
+        self.backend = backend
+        self._readers: dict[str, Any] = {}
+        torchvision.set_video_backend(backend)
+
+    def get(self, path: str) -> Any:
+        reader = self._readers.get(path)
+        if reader is None:
+            reader = torchvision.io.VideoReader(path, "video")
+            self._readers[path] = reader
+        return reader
+
+    def decode_frames(self, path: str, timestamps: list[float], tolerance_s: float = 1e-4) -> torch.Tensor:
+        reader = self.get(path)
+        first_ts = min(timestamps)
+        last_ts = max(timestamps)
+        reader.seek(first_ts, keyframes_only=self.backend == "pyav")
+
+        loaded_frames: list[torch.Tensor] = []
+        loaded_ts: list[float] = []
+        for frame in reader:
+            current_ts = float(frame["pts"])
+            loaded_frames.append(frame["data"])
+            loaded_ts.append(current_ts)
+            if current_ts >= last_ts:
+                break
+
+        query_ts = torch.tensor(timestamps, dtype=torch.float32)
+        loaded_ts_tensor = torch.tensor(loaded_ts, dtype=torch.float32)
+        distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
+        min_dist, argmin = distances.min(dim=1)
+        if not torch.all(min_dist < tolerance_s):
+            raise RuntimeError(
+                f"Video timestamps are outside tolerance: query={query_ts.tolist()} "
+                f"loaded={loaded_ts_tensor.tolist()} path={path}"
+            )
+        return torch.stack([loaded_frames[i] for i in argmin]).float() / 255.0
+
+
+@dataclass
+class A2DOpenLoopSample:
+    index: int
+    episode_index: int
+    task: str
+    state_raw: torch.Tensor
+    action_raw: torch.Tensor
+    inputs: dict[str, torch.Tensor]
+
+
+class A2DOpenLoopDataset:
+    image_keys = [
+        "observation.images.head",
+        "observation.images.hand_left",
+        "observation.images.hand_right",
+    ]
+    state_keys = [
+        "observation.states.joint.position",
+        "observation.states.effector.position",
+    ]
+    action_keys = [
+        "actions.joint.position",
+        "actions.effector.position",
+    ]
+
+    def __init__(
+        self,
+        dataset_root: str | Path,
+        *,
+        config: InternVLAA1Config,
+        train_stats: dict[str, Any],
+        image_offsets: tuple[int, int] = (-15, 0),
+        tolerance_s: float = 1e-4,
+    ) -> None:
+        self.root = Path(dataset_root)
+        self.config = config
+        self.info = _load_json(self.root / "meta" / "info.json")
+        self.dataset_stats = _load_json(self.root / "meta" / "stats.json")
+        self.data_rows = _load_parquet_rows(self.root / "data" / "chunk-000" / "file-000.parquet")
+        self.episode_rows = _load_parquet_rows(self.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
+        self.task_rows = _load_parquet_rows(self.root / "meta" / "tasks.parquet")
+
+        self.state_stats = _stack_stats(train_stats, self.state_keys)
+        self.action_stats = _stack_stats(train_stats, self.action_keys)
+        self.image_offsets = image_offsets
+        self.tolerance_s = tolerance_s
+        self.video_reader = TorchvisionVideoReaderCache(backend="pyav")
+
+    @property
+    def num_episodes(self) -> int:
+        return len(self.episode_rows)
+
+    @property
+    def physical_action_dim(self) -> int:
+        return 16
+
+    def episode_start_indices(self, max_episodes: int | None = None) -> list[tuple[int, list[int]]]:
+        rows = self.episode_rows if max_episodes is None else self.episode_rows[:max_episodes]
+        result = []
+        for ep in rows:
+            start = int(ep["dataset_from_index"])
+            end = int(ep["dataset_to_index"])
+            result.append((int(ep["episode_index"]), list(range(start, end, self.config.chunk_size))))
+        return result
+
+    def _task_text(self, task_index: int) -> str:
+        return self.task_rows[task_index]["__index_level_0__"]
+
+    def _episode_for_index(self, row: dict[str, Any]) -> dict[str, Any]:
+        return self.episode_rows[int(row["episode_index"])]
+
+    def _state_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.state_keys[0]] + row[self.state_keys[1]], dtype=torch.float32)
+
+    def _action_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.action_keys[0]] + row[self.action_keys[1]], dtype=torch.float32)
+
+    def _query_rows(self, idx: int, deltas: list[int]) -> list[dict[str, Any]]:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        start = int(episode["dataset_from_index"])
+        end = int(episode["dataset_to_index"])
+        return [self.data_rows[_clamp_index(idx + delta, start, end)] for delta in deltas]
+
+    def _decode_camera_history(
+        self, episode: dict[str, Any], camera_key: str, rows: list[dict[str, Any]]
+    ) -> torch.Tensor:
+        timestamps = [float(r["timestamp"]) for r in rows]
+        shifted = [float(episode[f"videos/{camera_key}/from_timestamp"]) + ts for ts in timestamps]
+        chunk_idx = int(episode[f"videos/{camera_key}/chunk_index"])
+        file_idx = int(episode[f"videos/{camera_key}/file_index"])
+        path = self.root / self.info["video_path"].format(
+            video_key=camera_key,
+            chunk_index=chunk_idx,
+            file_index=file_idx,
+        )
+        frames = self.video_reader.decode_frames(str(path), shifted, tolerance_s=self.tolerance_s)
+        return frames
+
+    def get_sample(self, idx: int) -> A2DOpenLoopSample:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        image_rows = self._query_rows(idx, list((-15, 0)))
+        action_rows = self._query_rows(idx, list(range(self.config.chunk_size)))
+        camera_images = [self._decode_camera_history(episode, camera_key, image_rows) for camera_key in self.image_keys]
+        state_raw = self._state_vector(row)
+        state_norm = normalize_vector(state_raw, self.state_stats)
+        action_raw = torch.stack([self._action_vector(action_row) for action_row in action_rows], dim=0)
+        task = self._task_text(int(row["task_index"]))
+        inputs = {
+            OBS_STATE: state_norm,
+            OBS_TASK: task,
+            f"{OBS_IMAGES}.image0": camera_images[0],
+            f"{OBS_IMAGES}.image1": camera_images[1],
+            f"{OBS_IMAGES}.image2": camera_images[2],
+            f"{OBS_IMAGES}.image0_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image1_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image2_mask": torch.tensor(True),
+        }
+        return A2DOpenLoopSample(
+            index=idx,
+            episode_index=int(row["episode_index"]),
+            task=task,
+            state_raw=state_raw,
+            action_raw=action_raw,
+            inputs=inputs,
+        )
+
+
+def collate_open_loop_samples(
+    samples: list[A2DOpenLoopSample],
+    *,
+    device: str,
+    dtype: torch.dtype,
+) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+    first = samples[0]
+    batch_inputs: dict[str, torch.Tensor] = {}
+    for key in first.inputs:
+        values = [sample.inputs[key] for sample in samples]
+        if isinstance(values[0], torch.Tensor):
+            tensor = torch.stack(values, dim=0)
+            if tensor.dtype in (torch.int64, torch.bool):
+                batch_inputs[key] = tensor.to(device=device)
+            else:
+                batch_inputs[key] = tensor.to(device=device, dtype=dtype)
+        else:
+            batch_inputs[key] = values
+
+    metadata = {
+        "indices": [sample.index for sample in samples],
+        "episode_indices": [sample.episode_index for sample in samples],
+        "tasks": [sample.task for sample in samples],
+        "state_raw": torch.stack([sample.state_raw for sample in samples], dim=0),
+        "action_raw": torch.stack([sample.action_raw for sample in samples], dim=0),
+    }
+    return batch_inputs, metadata
+
+
+def plot_prediction_series(
+    *,
+    series: dict[str, np.ndarray],
+    output_path: Path,
+    title: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    first_series = next(iter(series.values()))
+    num_dims = int(first_series.shape[1])
+    num_cols = 2
+    num_rows = math.ceil(num_dims / num_cols)
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(16, max(6, 3 * num_rows)))
+    axs = np.array(axs).reshape(-1)
+    x_values = np.arange(first_series.shape[0])
+    styles = [
+        ("Ground Truth", "blue", "-", 1.5),
+        ("VLLM", "green", "-.", 1.4),
+        ("Predicted", "red", "--", 1.4),
+    ]
+    style_map = {name: (color, linestyle, linewidth) for name, color, linestyle, linewidth in styles}
+
+    for dim in range(num_dims):
+        ax = axs[dim]
+        for name, values in series.items():
+            color, linestyle, linewidth = style_map.get(name, (None, "-", 1.2))
+            ax.plot(
+                x_values,
+                values[:, dim],
+                label=name,
+                color=color,
+                linestyle=linestyle,
+                linewidth=linewidth,
+            )
+        ax.set_title(f"Dimension {dim + 1}")
+        ax.set_xlabel("Time Step / Sample Index")
+        ax.set_ylabel(f"Value Dim {dim + 1}")
+        ax.legend(loc="upper right")
+        ax.grid(True, linestyle="--", alpha=0.7)
+
+    for dim in range(num_dims, len(axs)):
+        axs[dim].axis("off")
+
+    fig.suptitle(title, fontsize=16)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path)
+    plt.close(fig)
+
+
+def summarize_prediction_metrics(
+    *,
+    gt_tensor: torch.Tensor,
+    pred_tensor: torch.Tensor,
+    joint_dims: int = 14,
+) -> dict[str, float]:
+    return {
+        "mse": float(F.mse_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mae": float(F.l1_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mse_joint": float(F.mse_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mae_joint": float(F.l1_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mse_gripper": float(
+            F.mse_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+        "mae_gripper": float(
+            F.l1_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+    }
+
+
+def summarize_metric_list(values: list[float]) -> float | None:
+    return float(np.mean(values)) if values else None
+
+
+def run_open_loop_evaluation(
+    *,
+    mode: str,
+    policy,
+    config,
+    dataset,
+    train_meta,
+    collate_samples: Callable[..., tuple[dict[str, torch.Tensor], dict[str, Any]]],
+    run_sample_actions: Callable[[Any, dict[str, torch.Tensor], torch.Tensor], torch.Tensor],
+    num_episodes: int,
+    seed: int,
+    device: str,
+    dtype: torch.dtype,
+    output_dir: str | Path,
+    skip_plots: bool,
+) -> dict[str, Any]:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    plots_dir = output_dir / "plots"
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    mse_values: list[float] = []
+    mae_values: list[float] = []
+    mse_joint_values: list[float] = []
+    mae_joint_values: list[float] = []
+    mse_gripper_values: list[float] = []
+    mae_gripper_values: list[float] = []
+    per_episode: list[dict[str, Any]] = []
+
+    plotting_available = True
+    if not skip_plots:
+        try:
+            import matplotlib.pyplot as _plt  # noqa: F401
+        except ModuleNotFoundError:
+            plotting_available = False
+            print("[warning] matplotlib is not installed; plots will be skipped.")
+
+    episode_specs = dataset.episode_start_indices(max_episodes=num_episodes)
+    for episode_index, indices in episode_specs:
+        print(f"[{mode}] episode: {episode_index}")
+        pred_chunks = []
+        gt_chunks = []
+        visited_indices = []
+        task = None
+
+        for index in indices:
+            sample = dataset.get_sample(index)
+            task = sample.task
+            batch_inputs, meta = collate_samples([sample], device=device, dtype=dtype)
+            noise = make_shared_noise(
+                seed,
+                index,
+                (batch_inputs[OBS_STATE].shape[0], config.chunk_size, config.max_action_dim),
+                device,
+            )
+            with torch.no_grad():
+                pred = run_sample_actions(policy, batch_inputs, noise)
+
+            pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_phys = unnormalize_vector(pred, dataset.action_stats)
+            if train_meta.action_mode == "delta":
+                pred_phys[:, :, :14] += meta["state_raw"][:, None, :14]
+
+            gt_phys = meta["action_raw"][:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_chunks.append(pred_phys[0])
+            gt_chunks.append(gt_phys[0])
+            visited_indices.append(index)
+
+        pred_tensor = torch.cat(pred_chunks, dim=0)
+        gt_tensor = torch.cat(gt_chunks, dim=0)
+        metrics = summarize_prediction_metrics(gt_tensor=gt_tensor, pred_tensor=pred_tensor)
+
+        mse_values.append(metrics["mse"])
+        mae_values.append(metrics["mae"])
+        mse_joint_values.append(metrics["mse_joint"])
+        mae_joint_values.append(metrics["mae_joint"])
+        mse_gripper_values.append(metrics["mse_gripper"])
+        mae_gripper_values.append(metrics["mae_gripper"])
+
+        if not skip_plots and plotting_available:
+            plot_prediction_series(
+                series={
+                    "Ground Truth": gt_tensor.numpy(),
+                    "VLLM": pred_tensor.numpy(),
+                },
+                output_path=plots_dir / f"{mode}_open_loop_ep{episode_index}.jpg",
+                title=f"{mode} Ground Truth vs VLLM",
+            )
+
+        episode_log = {
+            "episode_id": int(episode_index),
+            "task": task,
+            "visited_indices": visited_indices,
+            "num_pred_steps": int(pred_tensor.shape[0]),
+            **metrics,
+        }
+        per_episode.append(episode_log)
+        print(json.dumps({"mode": mode, **episode_log}, ensure_ascii=False))
+
+    summary = {
+        "mode": mode,
+        "num_episodes": len(per_episode),
+        "mse": mse_values,
+        "mae": mae_values,
+        "average_mse": summarize_metric_list(mse_values),
+        "average_mae": summarize_metric_list(mae_values),
+        "mse_joint": mse_joint_values,
+        "mae_joint": mae_joint_values,
+        "average_mse_joint": summarize_metric_list(mse_joint_values),
+        "average_mae_joint": summarize_metric_list(mae_joint_values),
+        "mse_gripper": mse_gripper_values,
+        "mae_gripper": mae_gripper_values,
+        "average_mse_gripper": summarize_metric_list(mse_gripper_values),
+        "average_mae_gripper": summarize_metric_list(mae_gripper_values),
+        "episodes": per_episode,
+    }
+    with open(output_dir / "log.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2, ensure_ascii=False)
+    return summary

examples/offline_inference/internvla_a1/run.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+ : "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+ : "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+ : "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+ : "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+INTERNVLA_A1_OUTPUT_DIR="${INTERNVLA_A1_OUTPUT_DIR:-$REPO_ROOT/outputs/internvla_a1/vllm_infer}"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --output-dir "$INTERNVLA_A1_OUTPUT_DIR" \
+  --num-episodes "${INTERNVLA_A1_NUM_EPISODES:-1}" \
+  --attn-implementation eager \
+  "$@"

examples/offline_inference/magi_human/end2end.py

@@ -24,9 +24,6 @@ def parse_args():
     parser.add_argument("--width", type=int, default=448, help="Video width.")
     parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
     parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/moss_tts_nano/README.md

@@ -0,0 +1,101 @@
+# MOSS-TTS-Nano Offline Inference
+
+## Overview
+
+Single-stage offline TTS pipeline using the 0.1B MOSS-TTS-Nano AR LM and MOSS-Audio-Tokenizer-Nano codec. Outputs 48 kHz stereo WAV.
+
+## Quick Start
+
+```bash
+python end2end.py --text "Hello, this is MOSS-TTS-Nano."
+```
+
+The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face.
+
+## Usage
+
+```
+python end2end.py [OPTIONS]
+
+Options:
+  --text TEXT               Text to synthesize (default: "Hello, this is MOSS-TTS-Nano speaking.")
+  --voice VOICE             Built-in voice preset (default: Junhao). See voice table below.
+  --mode MODE               voice_clone (default) or continuation
+  --prompt-audio PATH       Reference WAV/MP3 for custom voice cloning
+  --prompt-text TEXT        Reference transcript (continuation mode)
+  --max-new-frames N        Max AR frames, default 375 (~14 s audio)
+  --seed INT                Random seed for reproducibility
+  --audio-temperature F     Audio sampling temperature (default: 0.8)
+  --audio-top-k N           Audio top-k sampling (default: 25)
+  --audio-top-p F           Audio top-p sampling (default: 0.95)
+  --text-temperature F      Text layer temperature (default: 1.0)
+  --batch                   Run a built-in batch of diverse samples (ZH/EN/FR)
+  --output-dir DIR          Directory for WAV outputs (default: /tmp/moss_tts_nano_output)
+  --deploy-config PATH      Override deploy YAML (defaults to vllm_omni/deploy/moss_tts_nano.yaml)
+  --stage-init-timeout INT  Timeout in seconds for stage init (default: 120)
+```
+
+## Examples
+
+```bash
+# Built-in Chinese voice
+python end2end.py --text "你好，这是MOSS-TTS-Nano的语音合成演示。" --voice Junhao
+
+# Built-in English voice
+python end2end.py --text "Hello from MOSS-TTS-Nano." --voice Ava
+
+# Custom voice clone
+python end2end.py \
+    --text "Hello, this is a cloned voice." \
+    --prompt-audio /path/to/reference.wav \
+    --prompt-text "Exact transcript of the reference audio."
+
+# Batch synthesis (ZH/EN/FR)
+python end2end.py --batch --output-dir /tmp/batch_output
+
+# Reproducible output
+python end2end.py --text "Deterministic test." --seed 42
+```
+
+## Built-in Voice Presets
+
+| Voice | Language |
+|-------|----------|
+| `Junhao` | ZH |
+| `Zhiming` | ZH |
+| `Weiguo` | ZH |
+| `Xiaoyu` | ZH |
+| `Yuewen` | ZH |
+| `Lingyu` | ZH |
+| `Ava` | EN |
+| `Bella` | EN |
+| `Adam` | EN |
+| `Nathan` | EN |
+| `Sakura` | JA |
+| `Yui` | JA |
+| `Aoi` | JA |
+| `Hina` | JA |
+| `Mei` | JA |
+
+## Deploy Config
+
+Runtime knobs live in `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded;
+override with `--deploy-config PATH`). Key stage-level settings:
+
+```yaml
+stages:
+  - stage_id: 0
+    gpu_memory_utilization: 0.3   # ~2 GB VRAM; increase for faster init
+    max_num_seqs: 4               # concurrent requests
+    max_model_len: 4096
+```
+
+## Output Format
+
+WAV files, 48 kHz, stereo (2-channel). The codec interleaves stereo as `[L, R, L, R, ...]` in the flat tensor returned by the model.
+
+## Troubleshooting
+
+- **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ torchcodec backend requires NVRTC. The model patches `torchaudio.load/save` automatically at load time to fall back to soundfile.
+- **`flash_attn not installed`**: The model falls back to `sdpa` attention automatically.
+- **Empty audio**: Check that `--text` is non-empty and the model loaded successfully (look for "MOSS-TTS-Nano LM loaded" in logs).

examples/offline_inference/moss_tts_nano/end2end.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference example for MOSS-TTS-Nano via vLLM-Omni.
+
+Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec
+both run inside one generation stage.  Output is 48 kHz stereo WAV.
+
+Supports:
+  - Built-in voice presets (--voice)
+  - Custom voice cloning via reference audio (--prompt-audio)
+  - Batch synthesis of multiple texts
+  - Reproducible output via --seed
+
+Usage examples:
+  # Built-in voice
+  python end2end.py --text "Hello from MOSS-TTS-Nano."
+
+  # Voice clone with reference audio
+  python end2end.py --text "Hello!" --prompt-audio /path/to/ref.wav
+
+  # Batch with different voices
+  python end2end.py --batch --output-dir /tmp/moss_output
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+from vllm import SamplingParams
+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")
+
+from vllm_omni import Omni  # noqa: E402
+
+MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
+
+BATCH_SAMPLES = [
+    {"text": "Hello, this is a test of MOSS-TTS-Nano.", "voice": "Ava", "label": "en_hello"},
+    {"text": "你好，这是 MOSS-TTS-Nano 的语音合成测试。", "voice": "Junhao", "label": "zh_hello"},
+    {"text": "Bonjour, ceci est un test de synthèse vocale.", "voice": "Bella", "label": "fr_hello"},
+]
+
+
+def build_request(
+    text: str,
+    voice: str = "Junhao",
+    mode: str = "voice_clone",
+    prompt_audio_path: str | None = None,
+    prompt_text: str | None = None,
+    max_new_frames: int = 375,
+    seed: int | None = None,
+    audio_temperature: float = 0.8,
+    audio_top_k: int = 25,
+    audio_top_p: float = 0.95,
+    text_temperature: float = 1.0,
+) -> dict:
+    """Build an Omni request payload for MOSS-TTS-Nano."""
+    additional: dict = {
+        "text": [text],
+        "voice": [voice],
+        "mode": [mode],
+        "max_new_frames": [max_new_frames],
+        "audio_temperature": [audio_temperature],
+        "audio_top_k": [audio_top_k],
+        "audio_top_p": [audio_top_p],
+        "text_temperature": [text_temperature],
+    }
+    if prompt_audio_path:
+        additional["prompt_audio_path"] = [str(prompt_audio_path)]
+    if prompt_text:
+        additional["prompt_text"] = [prompt_text]
+    if seed is not None:
+        additional["seed"] = [seed]
+
+    return {
+        "prompt": "<|im_start|>assistant\n",  # minimal placeholder prompt
+        "additional_information": additional,
+    }
+
+
+def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 48000) -> None:
+    audio_np = waveform.float().numpy()
+    # Reshape stereo: inference_stream yields interleaved [samples*2]; reshape to [samples, 2]
+    if audio_np.ndim == 1 and audio_np.shape[0] % 2 == 0:
+        audio_np = audio_np.reshape(-1, 2)
+    sf.write(path, audio_np, sample_rate)
+    print(f"  Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
+
+
+def main(args) -> None:
+    omni = Omni(
+        model=MODEL,
+        deploy_config=args.deploy_config,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    sampling_params = SamplingParams(
+        temperature=1.0,
+        top_p=1.0,
+        top_k=50,
+        max_tokens=4096,
+        seed=args.seed if args.seed is not None else 42,
+        detokenize=False,
+    )
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.batch:
+        print(f"Running batch synthesis ({len(BATCH_SAMPLES)} samples)...")
+        inputs = [build_request(s["text"], voice=s["voice"], seed=args.seed) for s in BATCH_SAMPLES]
+        params_list = [sampling_params] * len(inputs)
+    else:
+        print(f"Synthesizing: {args.text!r}")
+        inputs = build_request(
+            text=args.text,
+            voice=args.voice,
+            mode=args.mode,
+            prompt_audio_path=args.prompt_audio,
+            prompt_text=args.prompt_text,
+            max_new_frames=args.max_new_frames,
+            seed=args.seed,
+            audio_temperature=args.audio_temperature,
+            audio_top_k=args.audio_top_k,
+            audio_top_p=args.audio_top_p,
+            text_temperature=args.text_temperature,
+        )
+        params_list = sampling_params
+
+    for stage_outputs in omni.generate(inputs, params_list):
+        for i, req_output in enumerate(stage_outputs.request_output):
+            for j, out in enumerate(req_output.outputs):
+                mm = out.multimodal_output
+                if mm is None:
+                    print(f"  [req {i}] No audio output.")
+                    continue
+                audio = mm.get("audio")
+                sr_tensor = mm.get("sr")
+                if audio is None:
+                    print(f"  [req {i}] No waveform in multimodal_output.")
+                    continue
+                sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
+                label = BATCH_SAMPLES[i]["label"] if args.batch else f"output_{i}_{j}"
+                out_path = str(output_dir / f"{label}.wav")
+                save_audio(audio.cpu(), out_path, sr)
+
+    print("Done.")
+
+
+def parse_args():
+    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("--voice", default="Junhao", help="Built-in voice preset name.")
+    parser.add_argument("--mode", default="voice_clone", choices=["voice_clone", "continuation"])
+    parser.add_argument("--prompt-audio", default=None, help="Path to reference audio for voice cloning.")
+    parser.add_argument("--prompt-text", default=None, help="Reference transcript (continuation mode).")
+    parser.add_argument("--max-new-frames", type=int, default=375, help="Max AR frames (~14s at default).")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed.")
+    parser.add_argument("--audio-temperature", type=float, default=0.8)
+    parser.add_argument("--audio-top-k", type=int, default=25)
+    parser.add_argument("--audio-top-p", type=float, default=0.95)
+    parser.add_argument("--text-temperature", type=float, default=1.0)
+    parser.add_argument("--batch", action="store_true", help="Run built-in batch of diverse samples.")
+    parser.add_argument("--output-dir", default="/tmp/moss_tts_nano_output", help="Directory for WAV outputs.")
+    parser.add_argument(
+        "--deploy-config",
+        default=None,
+        help="Path to a deploy YAML; leave unset to auto-load vllm_omni/deploy/moss_tts_nano.yaml.",
+    )
+    parser.add_argument("--stage-init-timeout", type=int, default=120)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -242,8 +242,7 @@ async def run_single_request(
                 if stage_0_first_output_ts is None:
                     stage_0_first_output_ts = time.perf_counter()
                 text_output = output.outputs[0].text
-                if output.finished:
-                    text_parts.append(text_output)
+                text_parts.append(text_output)
             elif omni_output.final_output_type == "audio":
                 mm_out = output.outputs[0].multimodal_output
                 if mm_out and "audio" in mm_out:
@@ -293,7 +292,7 @@ async def run_single_request(
     if text_parts:
         text_file = os.path.join(output_dir, f"{request_id}.txt")
         with open(text_file, "w", encoding="utf-8") as f:
-            f.write("\n".join(text_parts))
+            f.write("".join(text_parts))
         result["saved_files"].append(text_file)
         print(
             f"[Request {request_id}] Text saved to {text_file} "
@@ -393,6 +392,10 @@ async def run_all(args):
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
+        #
+        # NOTE: Since we do not set the sampling params directly, .generate in
+        # will automatically set the output kind to delta, since this is what
+        # makes sense for most multimodal use-cases.
         sampling_params_list = None
 
         output_dir = args.output_dir

examples/offline_inference/qwen3_tts/end2end.py

@@ -83,12 +83,29 @@ def _estimate_prompt_len(
             if not isinstance(audio_path, str) or not audio_path.strip():
                 return None
             try:
-                from vllm.multimodal.media import MediaConnector
+                from urllib.parse import urlparse
 
-                connector = MediaConnector(allowed_local_media_path="/")
-                audio, sr = connector.fetch_audio(audio_path)
                 import numpy as np
 
+                def _is_url(path: str) -> bool:
+                    try:
+                        parsed = urlparse(path)
+                        if parsed.scheme in ("http", "https"):
+                            return bool(parsed.netloc)
+                        return parsed.scheme in ("file", "data")
+                    except Exception:
+                        return False
+
+                if _is_url(audio_path):
+                    from vllm.multimodal.media import MediaConnector
+
+                    connector = MediaConnector(allowed_local_media_path="/")
+                    audio, sr = connector.fetch_audio(audio_path)
+                else:
+                    from vllm.multimodal.media.audio import load_audio
+
+                    audio, sr = load_audio(audio_path, sr=None, mono=True)
+
                 wav_np = np.asarray(audio, dtype=np.float32)
 
                 if speech_tok is not None:

examples/offline_inference/text_to_image/text_to_image.py

@@ -9,7 +9,7 @@ from typing import Any
 
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
+from vllm_omni.diffusion.data import logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
@@ -342,17 +342,6 @@ def main():
             #       (e.g., QwenImagePipeline or FluxPipeline)
         }
 
-    # assert args.ring_degree == 1, "Ring attention is not supported yet"
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        ulysses_mode=args.ulysses_mode,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        enable_expert_parallel=args.enable_expert_parallel,
-    )
-
     profiler_enabled = args.profiler_config is not None
 
     # Prepare LoRA kwargs for Omni initialization
@@ -388,7 +377,13 @@ def main():
         "cache_backend": args.cache_backend,
         "cache_config": cache_config,
         "enable_cache_dit_summary": args.enable_cache_dit_summary,
-        "parallel_config": parallel_config,
+        "ulysses_degree": args.ulysses_degree,
+        "ring_degree": args.ring_degree,
+        "ulysses_mode": args.ulysses_mode,
+        "cfg_parallel_size": args.cfg_parallel_size,
+        "tensor_parallel_size": args.tensor_parallel_size,
+        "vae_patch_parallel_size": args.vae_patch_parallel_size,
+        "enable_expert_parallel": args.enable_expert_parallel,
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
         "mode": "text-to-image",

examples/offline_inference/text_to_video/text_to_video.py

@@ -226,9 +226,6 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of HSDP replica groups.",
     )
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/vace/vace_video_generation.py

@@ -71,9 +71,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
     parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
     parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/online_serving/moss_tts_nano/README.md

@@ -0,0 +1,158 @@
+# MOSS-TTS-Nano
+
+## Model checkpoint
+
+| Model | Description |
+|-------|-------------|
+| `OpenMOSS-Team/MOSS-TTS-Nano` | 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec, 48 kHz stereo, ZH/EN/JA |
+
+## Gradio Demo
+
+An interactive Gradio demo is available with multilingual voice presets, custom voice cloning, and streaming support.
+
+```bash
+# Option 1: Launch server + Gradio together
+./run_gradio_demo.sh
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8091
+```
+
+Then open http://localhost:7860 in your browser.
+
+Features:
+
+- 15 built-in voice presets (6 ZH · 4 EN · 5 JA)
+- Custom voice cloning from uploaded reference audio
+- Streaming mode (progressive PCM playback)
+
+## Launch the Server
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
+```
+
+The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no
+`--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags
+are needed.
+
+Or use the convenience script:
+
+```bash
+./run_server.sh
+```
+
+## Send TTS Request
+
+### Using curl
+
+```bash
+# Built-in voice preset
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "你好，这是语音合成测试。",
+        "voice": "Junhao",
+        "response_format": "wav"
+    }' --output output.wav
+
+# English preset
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is MOSS-TTS-Nano.",
+        "voice": "Ava",
+        "response_format": "wav"
+    }' --output output_en.wav
+```
+
+### Custom Voice Cloning
+
+Provide a reference audio (base64 data URL) and its transcript:
+
+```bash
+REF_AUDIO=$(base64 -w 0 /path/to/reference.wav)
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d "{
+        \"input\": \"Hello, this is a cloned voice.\",
+        \"voice\": \"Junhao\",
+        \"ref_audio\": \"data:audio/wav;base64,${REF_AUDIO}\",
+        \"ref_text\": \"Exact transcript of the reference audio.\"
+    }" --output cloned.wav
+```
+
+### Using Python
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "你好，这是语音合成测试。",
+        "voice": "Junhao",
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### Streaming
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, streaming output from MOSS-TTS-Nano.",
+        "voice": "Ava",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 2 -
+```
+
+**Note:** MOSS-TTS-Nano outputs at 48 kHz stereo (2-channel).
+
+## API Parameters
+
+MOSS-TTS-Nano uses the standard `/v1/audio/speech` endpoint.
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text to synthesize (ZH / EN / JA) |
+| `voice` | string | `"Junhao"` | Built-in voice name (see table below) |
+| `response_format` | string | `"wav"` | Audio format: wav, mp3, flac, pcm |
+| `ref_audio` | string | null | Base64 data URL for custom voice cloning |
+| `ref_text` | string | null | Transcript of reference audio (required when `ref_audio` is set) |
+| `stream` | bool | false | Stream raw PCM chunks |
+| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
+
+## Built-in Voice Presets
+
+| Voice | Language | Description |
+|-------|----------|-------------|
+| `Junhao` | ZH | Male, standard Mandarin |
+| `Zhiming` | ZH | Male |
+| `Weiguo` | ZH | Male |
+| `Xiaoyu` | ZH | Female |
+| `Yuewen` | ZH | Female |
+| `Lingyu` | ZH | Female |
+| `Ava` | EN | Female, American English |
+| `Bella` | EN | Female |
+| `Adam` | EN | Male |
+| `Nathan` | EN | Male |
+| `Sakura` | JA | Female |
+| `Yui` | JA | Female |
+| `Aoi` | JA | Female |
+| `Hina` | JA | Female |
+| `Mei` | JA | Female |
+
+## Troubleshooting
+
+1. **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ defaults to torchcodec which requires NVRTC. vLLM-Omni patches this automatically at model load time to use soundfile instead.
+2. **Connection refused**: Ensure the server is running on the correct port.
+3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings.
+4. **Out of memory**: The default `gpu_memory_utilization=0.3` is conservative. Increase it in the stage config if you have more VRAM available.

examples/online_serving/moss_tts_nano/gradio_demo.py

@@ -0,0 +1,690 @@
+"""Gradio demo for MOSS-TTS-Nano with gapless streaming audio playback.
+
+Uses a custom AudioWorklet-based player (adapted from the Qwen3-TTS demo)
+for gap-free streaming.  Audio is streamed from the vLLM-Omni server
+through a same-origin proxy and played via the Web Audio API's AudioWorklet.
+
+MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved L,R,L,R,...).
+The AudioWorklet downmixes stereo to mono during playback.
+
+Also supports non-streaming mode (full audio download) via gr.Audio.
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8091
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import json
+import logging
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+
+import httpx
+import numpy as np
+import soundfile as sf
+from fastapi import FastAPI, Request
+from fastapi.responses import Response, StreamingResponse
+
+logger = logging.getLogger(__name__)
+
+# MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved).
+PCM_SAMPLE_RATE = 48000
+PCM_CHANNELS = 2
+DEFAULT_API_BASE = "http://localhost:8091"
+
+# Built-in voice presets (matches the MOSS-TTS-Nano GitHub repo assets/)
+VOICE_PRESETS = {
+    # Chinese
+    "Junhao (ZH)": "Junhao",
+    "Zhiming (ZH)": "Zhiming",
+    "Weiguo (ZH)": "Weiguo",
+    "Xiaoyu (ZH)": "Xiaoyu",
+    "Yuewen (ZH)": "Yuewen",
+    "Lingyu (ZH)": "Lingyu",
+    # English
+    "Ava (EN)": "Ava",
+    "Bella (EN)": "Bella",
+    "Adam (EN)": "Adam",
+    "Nathan (EN)": "Nathan",
+    # Japanese
+    "Sakura (JA)": "Sakura",
+    "Yui (JA)": "Yui",
+    "Aoi (JA)": "Aoi",
+    "Hina (JA)": "Hina",
+    "Mei (JA)": "Mei",
+}
+PRESET_LABELS = list(VOICE_PRESETS.keys())
+DEFAULT_PRESET = "Junhao (ZH)"
+
+# ── AudioWorklet processor (loaded in browser via Blob URL) ──────────
+# Downmixes stereo-interleaved int16 PCM to mono float32 for playback.
+WORKLET_JS = r"""
+class TTSPlaybackProcessor extends AudioWorkletProcessor {
+    constructor() {
+        super();
+        this.queue = [];
+        this.buf = null;
+        this.pos = 0;
+        this.playing = false;
+        this.played = 0;
+        this.port.onmessage = (e) => {
+            if (e.data && e.data.type === 'clear') {
+                this.queue = []; this.buf = null; this.pos = 0; this.played = 0;
+                if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); }
+                return;
+            }
+            // Input: Int16Array of stereo-interleaved samples [L0,R0,L1,R1,...]
+            // Convert to mono float32 by averaging L+R channels.
+            const pcm = e.data;
+            const monoLen = Math.floor(pcm.length / 2);
+            const mono = new Float32Array(monoLen);
+            for (let i = 0; i < monoLen; i++) {
+                mono[i] = (pcm[i * 2] + pcm[i * 2 + 1]) / 65536.0;
+            }
+            this.queue.push(mono);
+        };
+    }
+    process(inputs, outputs) {
+        const out = outputs[0][0];
+        for (let i = 0; i < out.length; i++) {
+            if (!this.buf || this.pos >= this.buf.length) {
+                if (this.queue.length > 0) {
+                    this.buf = this.queue.shift(); this.pos = 0;
+                } else {
+                    for (let j = i; j < out.length; j++) out[j] = 0;
+                    if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); }
+                    return true;
+                }
+            }
+            out[i] = this.buf[this.pos++];
+            this.played++;
+        }
+        if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); }
+        return true;
+    }
+}
+registerProcessor('tts-playback-processor', TTSPlaybackProcessor);
+"""
+
+# ── Player HTML (container with metric cards) ────────────────────────
+PLAYER_HTML = """
+<div id="tts-player">
+  <div style="display:flex; align-items:center; gap:10px;">
+    <div id="tts-status-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
+    <span id="tts-status" style="font-weight:600;font-size:1.05em;">Ready</span>
+    <button id="tts-stop-btn" onclick="window.ttsStop()"
+      style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
+             background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em; transition:all 0.15s;">Stop</button>
+  </div>
+  <div id="tts-metrics" style="display:none; grid-template-columns:repeat(4,1fr); gap:10px; margin-top:12px;">
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">TTFP</div>
+      <div id="tts-m-ttfp" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">RTF</div>
+      <div id="tts-m-rtf" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Audio</div>
+      <div id="tts-m-dur" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Speed</div>
+      <div id="tts-m-speed" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+  </div>
+  <div id="tts-rtf-bar-wrap" style="display:none; background:#e8ecf1; border-radius:4px; height:20px; overflow:hidden; position:relative; margin-top:10px;">
+    <div id="tts-rtf-bar" style="height:100%; border-radius:4px; transition:width 0.3s ease, background 0.3s ease; width:0%;"></div>
+    <span id="tts-rtf-label" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);font-size:0.75em;font-weight:600;color:#444;"></span>
+  </div>
+  <div id="tts-elapsed" style="display:none; margin-top:6px; font-size:0.8em; color:#999; text-align:right;"></div>
+</div>
+"""
+
+
+def _build_player_js(sample_rate: int) -> str:
+    """Build the JavaScript that powers the AudioWorklet player."""
+    return f"""
+    <script>
+    const SR = {sample_rate};
+    const WC = {json.dumps(WORKLET_JS)};
+    let ctx = null, node = null, abort = null, gen = false, st = {{}};
+
+    async function init() {{
+        if (ctx) return;
+        ctx = new AudioContext({{ sampleRate: SR }});
+        const b = new Blob([WC], {{ type: 'application/javascript' }});
+        const u = URL.createObjectURL(b);
+        await ctx.audioWorklet.addModule(u);
+        URL.revokeObjectURL(u);
+        node = new AudioWorkletNode(ctx, 'tts-playback-processor');
+        node.connect(ctx.destination);
+        node.port.onmessage = (e) => {{
+            if (e.data.type === 'started') setStatus('Playing...', '#2E7D32');
+            else if (e.data.type === 'stopped' && !gen) {{
+                setStatus('Done', '#2E7D32'); showStats(true);
+                const btn = document.getElementById('tts-stop-btn');
+                if (btn) btn.style.display = 'none';
+            }}
+        }};
+    }}
+
+    function setStatus(text, color) {{
+        const s = document.getElementById('tts-status');
+        const d = document.getElementById('tts-status-dot');
+        if (s) s.textContent = text;
+        if (d) d.style.background = color || '#ccc';
+    }}
+
+    function showStats(fin) {{
+        if (!st.t0) return;
+        const elapsed = (fin && st.streamEnd ? (st.streamEnd - st.t0) : (performance.now() - st.t0)) / 1000;
+        // st.samples counts stereo int16 samples; mono frames = samples / 2
+        const dur = (st.samples / 2) / SR;
+        const mTtfp = document.getElementById('tts-m-ttfp');
+        const mRtf = document.getElementById('tts-m-rtf');
+        const mDur = document.getElementById('tts-m-dur');
+        const mSpeed = document.getElementById('tts-m-speed');
+        const bar = document.getElementById('tts-rtf-bar');
+        const barLabel = document.getElementById('tts-rtf-label');
+        const elapsedEl = document.getElementById('tts-elapsed');
+
+        if (mTtfp && st.ttfp != null) mTtfp.textContent = st.ttfp.toFixed(0) + 'ms';
+        if (mDur) mDur.textContent = dur.toFixed(1) + 's';
+
+        if (dur > 0 && elapsed > 0) {{
+            const rtf = elapsed / dur;
+            const speed = 1 / rtf;
+            if (mRtf) {{
+                mRtf.textContent = rtf.toFixed(2) + 'x';
+                mRtf.style.color = rtf < 1 ? '#2E7D32' : rtf < 1.5 ? '#e8a317' : '#EF5552';
+            }}
+            if (mSpeed) {{
+                mSpeed.textContent = speed.toFixed(1) + 'x';
+                mSpeed.style.color = speed > 1 ? '#2E7D32' : speed > 0.7 ? '#e8a317' : '#EF5552';
+            }}
+            if (bar) {{
+                const pct = Math.min(speed / 3 * 100, 100);
+                bar.style.width = pct + '%';
+                bar.style.background = speed > 1 ? 'linear-gradient(90deg,#66BB6A,#2E7D32)' : speed > 0.7 ? 'linear-gradient(90deg,#e8a317,#f0b866)' : 'linear-gradient(90deg,#EF5552,#f87171)';
+            }}
+            if (barLabel) barLabel.textContent = speed.toFixed(1) + 'x realtime';
+        }}
+        if (elapsedEl) {{
+            elapsedEl.style.display = 'block';
+            elapsedEl.textContent = fin ? 'Completed in ' + elapsed.toFixed(1) + 's  (' + st.chunks + ' chunks)' : elapsed.toFixed(1) + 's elapsed  (' + st.chunks + ' chunks)';
+        }}
+    }}
+
+    window.ttsStop = function() {{
+        if (abort) abort.abort();
+        if (node) node.port.postMessage({{ type: 'clear' }});
+        gen = false;
+        setStatus('Stopped', '#999');
+        const btn = document.getElementById('tts-stop-btn');
+        if (btn) btn.style.display = 'none';
+    }};
+
+    window.ttsGenerate = async function(payload) {{
+        try {{ await init(); if (ctx.state === 'suspended') await ctx.resume(); }}
+        catch (e) {{ const s = document.getElementById('tts-status'); if (s) s.textContent = 'Audio init error: ' + e.message; return; }}
+
+        if (abort) abort.abort();
+        node.port.postMessage({{ type: 'clear' }});
+        await new Promise(r => setTimeout(r, 50));
+        node.port.postMessage({{ type: 'clear' }});
+
+        gen = true;
+        st = {{ t0: null, chunks: 0, samples: 0, ttfp: null }};
+        setStatus('Connecting...', '#2E7D32');
+        const bEl = document.getElementById('tts-stop-btn');
+        if (bEl) bEl.style.display = 'inline-block';
+        const mp = document.getElementById('tts-metrics');
+        if (mp) {{ mp.style.display = 'grid'; ['tts-m-ttfp','tts-m-rtf','tts-m-dur','tts-m-speed'].forEach(id => {{ const e = document.getElementById(id); if(e) {{ e.textContent = '\\u2014'; e.style.color = '#333'; }} }}); }}
+        const bw = document.getElementById('tts-rtf-bar-wrap');
+        if (bw) bw.style.display = 'block';
+        const bar = document.getElementById('tts-rtf-bar');
+        if (bar) bar.style.width = '0%';
+        const bl = document.getElementById('tts-rtf-label');
+        if (bl) bl.textContent = '';
+        const ee = document.getElementById('tts-elapsed');
+        if (ee) {{ ee.style.display = 'none'; ee.textContent = ''; }}
+        abort = new AbortController();
+
+        try {{
+            st.t0 = performance.now();
+            const r = await fetch('/proxy/v1/audio/speech', {{
+                method: 'POST',
+                headers: {{ 'Content-Type': 'application/json' }},
+                body: JSON.stringify(payload),
+                signal: abort.signal,
+            }});
+            if (!r.ok) {{ const t = await r.text(); throw new Error('Server ' + r.status + ': ' + t.slice(0, 200)); }}
+            setStatus('Streaming...', '#43A047');
+
+            const reader = r.body.getReader();
+            let left = new Uint8Array(0);
+            while (true) {{
+                const {{ done, value }} = await reader.read();
+                if (done) break;
+                let raw;
+                if (left.length > 0) {{
+                    raw = new Uint8Array(left.length + value.length);
+                    raw.set(left); raw.set(value, left.length);
+                }} else {{ raw = value; }}
+                const usable = raw.length - (raw.length % 2);
+                left = usable < raw.length ? raw.slice(usable) : new Uint8Array(0);
+                if (usable > 0) {{
+                    const ab = new ArrayBuffer(usable);
+                    new Uint8Array(ab).set(raw.subarray(0, usable));
+                    const pcm = new Int16Array(ab);
+                    node.port.postMessage(pcm);
+                    st.chunks++;
+                    st.samples += pcm.length;
+                    if (st.ttfp == null) st.ttfp = performance.now() - st.t0;
+                    showStats(false);
+                }}
+            }}
+        }} catch (e) {{
+            if (e.name !== 'AbortError') {{
+                setStatus('Error: ' + e.message, '#EF5552');
+                console.error('TTS error:', e);
+            }}
+        }} finally {{
+            st.streamEnd = performance.now();
+            showStats(true);
+            gen = false;
+            if (st.samples > 0) setStatus('Finishing playback...', '#2E7D32');
+            else {{
+                setStatus('No audio received', '#999');
+                if (bEl) bEl.style.display = 'none';
+            }}
+        }}
+    }};
+    </script>
+"""
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    voice_preset: str,
+    ref_audio: tuple | None,
+    ref_text: str,
+    response_format: str = "pcm",
+    stream: bool = True,
+) -> dict:
+    """Build the /v1/audio/speech request payload."""
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    voice_name = VOICE_PRESETS.get(voice_preset, "Junhao")
+
+    payload: dict = {
+        "input": text.strip(),
+        "voice": voice_name,
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+    }
+
+    has_ref = ref_audio is not None
+    if has_ref:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+        ref_text_stripped = ref_text.strip() if ref_text else ""
+        if ref_text_stripped:
+            payload["ref_text"] = ref_text_stripped
+        else:
+            raise gr.Error("Voice cloning requires a transcript of the reference audio.")
+
+    return payload
+
+
+def generate_speech(
+    api_base: str,
+    text: str,
+    voice_preset: str,
+    ref_audio: tuple | None,
+    ref_text: str,
+    response_format: str,
+) -> tuple:
+    """Non-streaming: call /v1/audio/speech and return full audio as (sr, np_array)."""
+    payload = build_payload(text, voice_preset, ref_audio, ref_text, response_format, stream=False)
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    try:
+        if response_format == "pcm":
+            # PCM is stereo interleaved; downmix to mono for Gradio.
+            samples = np.frombuffer(resp.content, dtype=np.int16)
+            n = len(samples) - (len(samples) % 2)
+            left = samples[:n:2].astype(np.float32)
+            right = samples[1:n:2].astype(np.float32)
+            mono = (left + right) / 65536.0
+            return (PCM_SAMPLE_RATE, mono)
+        # WAV/MP3/FLAC: soundfile handles header parsing.
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np.mean(axis=1)
+        return (int(sample_rate), audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}")
+
+
+def create_app(api_base: str):
+    """Create the FastAPI app with streaming proxy + Gradio UI."""
+    fastapi_app = FastAPI()
+
+    # ── Streaming proxy (same-origin, no CORS issues) ────────────
+    @fastapi_app.post("/proxy/v1/audio/speech")
+    async def proxy_speech(request: Request):
+        body = await request.json()
+        logger.info(
+            "Proxy request: input=%r voice=%s stream=%s",
+            body.get("input", "")[:50],
+            body.get("voice"),
+            body.get("stream"),
+        )
+        try:
+            client = httpx.AsyncClient(timeout=300)
+            resp = await client.send(
+                client.build_request(
+                    "POST",
+                    f"{api_base}/v1/audio/speech",
+                    json=body,
+                    headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"},
+                ),
+                stream=True,
+            )
+        except Exception as exc:
+            logger.exception("Proxy connection error")
+            await client.aclose()
+            return Response(content=str(exc), status_code=502)
+
+        if resp.status_code != 200:
+            content = await resp.aread()
+            logger.error("Proxy upstream error %d: %s", resp.status_code, content[:200])
+            await resp.aclose()
+            await client.aclose()
+            return Response(content=content, status_code=resp.status_code)
+
+        async def relay():
+            total = 0
+            try:
+                async for chunk in resp.aiter_bytes():
+                    total += len(chunk)
+                    yield chunk
+            except Exception:
+                logger.exception("Proxy relay error after %d bytes", total)
+            finally:
+                logger.info("Proxy relay done: %d bytes", total)
+                await resp.aclose()
+                await client.aclose()
+
+        return StreamingResponse(relay(), media_type="application/octet-stream")
+
+    # ── Gradio UI ────────────────────────────────────────────────
+    css = """
+    #generate-btn button { width: 100%; }
+    #streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; }
+    """
+
+    theme = gr.themes.Default(
+        primary_hue=gr.themes.Color(
+            c50="#E8F5E9",
+            c100="#C8E6C9",
+            c200="#A5D6A7",
+            c300="#81C784",
+            c400="#66BB6A",
+            c500="#4CAF50",
+            c600="#43A047",
+            c700="#388E3C",
+            c800="#2E7D32",
+            c900="#1B5E20",
+            c950="#0D3B10",
+        ),
+    )
+
+    with gr.Blocks(title="MOSS-TTS-Nano Demo") as demo:
+        gr.HTML(f"""
+        <div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
+          <img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
+               alt="vLLM-Omni" style="height:42px;">
+          <div>
+            <h1 style="margin:0; font-size:1.5em;">MOSS-TTS-Nano Streaming Demo</h1>
+            <span style="font-size:0.85em; color:#666;">
+              0.1B AR LM + MOSS-Audio-Tokenizer-Nano &middot; 48 kHz &middot; ZH / EN / JA
+              &nbsp;&middot;&nbsp; Served by
+              <a href="https://github.com/vllm-project/vllm-omni" target="_blank"
+              style="color:#43A047; text-decoration:none; font-weight:600;">vLLM-Omni</a>
+              &nbsp;&middot;&nbsp; <code style="background:#eef2f7; padding:2px 6px; border-radius:4px; font-size:0.9em;">{api_base}</code>
+            </span>
+          </div>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text in Chinese, English, or Japanese...",
+                    lines=4,
+                )
+
+                voice_preset = gr.Dropdown(
+                    choices=PRESET_LABELS,
+                    value=DEFAULT_PRESET,
+                    label="Voice Preset",
+                    info="Built-in multilingual voices. ZH = Chinese, EN = English, JA = Japanese.",
+                )
+
+                with gr.Accordion("Custom Voice Clone (optional)", open=False):
+                    gr.Markdown(
+                        "Upload a reference audio clip (10-30 s) and provide its transcript "
+                        "to clone that voice instead of using the built-in preset."
+                    )
+                    ref_audio = gr.Audio(
+                        label="Reference Audio",
+                        type="numpy",
+                        sources=["upload", "microphone"],
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"],
+                        value="pcm",
+                        label="Audio Format",
+                        interactive=False,
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream (gapless)",
+                        value=True,
+                        info="AudioWorklet streaming (recommended)",
+                        scale=1,
+                    )
+
+                with gr.Row():
+                    generate_btn = gr.Button(
+                        "Generate Speech",
+                        variant="primary",
+                        size="lg",
+                        elem_id="generate-btn",
+                        scale=3,
+                    )
+                    reset_btn = gr.Button(
+                        "Reset",
+                        variant="secondary",
+                        size="lg",
+                        scale=1,
+                    )
+
+            with gr.Column(scale=2):
+                player_html = gr.HTML(
+                    value=PLAYER_HTML,
+                    visible=True,
+                    label="streaming player",
+                    elem_id="streaming-player",
+                )
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    autoplay=True,
+                    visible=False,
+                )
+                gr.Examples(
+                    examples=[
+                        ["Hello, this is MOSS-TTS-Nano speaking.", "Ava (EN)"],
+                        ["The quick brown fox jumps over the lazy dog.", "Adam (EN)"],
+                        ["MOSS-TTS-Nanoは、軽量なテキスト読み上げモデルです。", "Sakura (JA)"],
+                    ],
+                    inputs=[text_input, voice_preset],
+                    label="Examples",
+                )
+                gr.HTML("""
+                <div style="text-align:center; padding:8px 0; margin-top:4px;">
+                  <a href="https://github.com/vllm-project/vllm-omni" target="_blank">
+                    <img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
+                         alt="vLLM-Omni" style="height:28px; opacity:0.7;">
+                  </a>
+                </div>
+                """)
+
+        # Hidden textbox to pass payload from Python -> JavaScript
+        hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload")
+
+        # ── Event wiring ─────────────────────────────────────────
+        def on_stream_change(stream: bool):
+            if stream:
+                return (
+                    gr.update(value="pcm", interactive=False),
+                    gr.update(visible=True),  # player
+                    gr.update(visible=False),  # audio
+                )
+            return (
+                gr.update(value="wav", interactive=True),
+                gr.update(visible=False),
+                gr.update(visible=True),
+            )
+
+        stream_checkbox.change(
+            fn=on_stream_change,
+            inputs=[stream_checkbox],
+            outputs=[response_format, player_html, audio_output],
+        )
+
+        def on_reset():
+            return (
+                "",  # text
+                None,  # audio_output
+                "",  # hidden_payload
+                PLAYER_HTML,  # reset player
+            )
+
+        reset_btn.click(
+            fn=on_reset,
+            outputs=[text_input, audio_output, hidden_payload, player_html],
+            js="() => { if (window.ttsStop) window.ttsStop(); }",
+        )
+
+        all_inputs = [text_input, voice_preset, ref_audio, ref_text, response_format]
+
+        def on_generate(stream_enabled, *args):
+            import time as _time
+
+            text, voice, ref_a, ref_t, _fmt = args
+            if stream_enabled:
+                payload = build_payload(text, voice, ref_a, ref_t, stream=True)
+                payload["_nonce"] = int(_time.time() * 1000)
+                return json.dumps(payload), gr.update()
+            else:
+                audio = generate_speech(api_base, text, voice, ref_a, ref_t, _fmt)
+                return "", audio
+
+        generate_btn.click(
+            fn=on_generate,
+            inputs=[stream_checkbox] + all_inputs,
+            outputs=[hidden_payload, audio_output],
+        ).then(
+            fn=lambda p: p,
+            inputs=[hidden_payload],
+            outputs=[hidden_payload],
+            js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }",
+        )
+
+        demo.queue()
+
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        css=css,
+        theme=theme,
+        head=_build_player_js(PCM_SAMPLE_RATE),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Gradio demo for MOSS-TTS-Nano with gapless AudioWorklet streaming.",
+    )
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
+    parser.add_argument("--share", action="store_true", help="Share publicly via Gradio tunnel")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+    print(f"Connecting to vLLM-Omni server at: {args.api_base}")
+
+    app = create_app(args.api_base)
+
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/moss_tts_nano/run_gradio_demo.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# Launch MOSS-TTS-Nano server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
+PORT="${PORT:-8091}"
+GRADIO_PORT="${GRADIO_PORT:-7860}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+echo "Starting MOSS-TTS-Nano server (port $PORT)..."
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --omni &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+# Wait for server to be ready.
+echo "Waiting for server to start..."
+for i in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"

examples/online_serving/moss_tts_nano/run_server.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+# Launch vLLM-Omni server for MOSS-TTS-Nano
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
+PORT="${PORT:-8091}"
+
+echo "Starting MOSS-TTS-Nano server with model: $MODEL"
+
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --omni

examples/online_serving/qwen3_omni/README.md

@@ -31,11 +31,13 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 ### Launch individual stages (stage-based CLI)
 
 Use the stage-based CLI when you want to run one stage per process.
+The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
+`CUDA_VISIBLE_DEVICES`.
 
 **1. Stage 0 (Thinker + API server)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -45,7 +47,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **2. Stage 1 (Talker)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -55,7 +57,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **3. Stage 2 (Code2Wav)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 2 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -78,7 +80,7 @@ for this category of optimization. Given that each instantiation strictly initia
 parameter flags can be systematically assigned directly onto that specific stage's command sequence:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

examples/online_serving/stable_audio/README.md

@@ -0,0 +1,234 @@
+# Stable Audio Online Serving
+
+Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
+
+## Features
+
+- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
+- **Flexible control**: Adjust audio length, guidance scale, inference steps
+- **Quality control**: Use negative prompts to avoid unwanted characteristics
+- **Reproducible**: Set random seed for deterministic generation
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### 2. Generate Audio
+
+#### Using curl
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+#### Using Python Client
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring" \
+    --audio_length 10.0 \
+    --output cat.wav
+```
+
+#### Using Bash Script
+
+```bash
+bash curl_examples.sh
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+```
+
+### Request Body
+
+```json
+{
+    "input": "Text description of the audio",
+    "audio_length": 10.0,
+    "audio_start": 0.0,
+    "negative_prompt": "Low quality",
+    "guidance_scale": 7.0,
+    "num_inference_steps": 100,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+### Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducibility |
+| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Usage Examples
+
+### Basic Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+## Python Client Examples
+
+### Simple Generation
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+```
+
+### Custom Parameters
+
+```bash
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+### Different Output Format
+
+```bash
+python stable_audio_client.py \
+    --text "Guitar playing" \
+    --response_format mp3 \
+    --output guitar.mp3
+```
+
+## Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality
+   - 100 steps: Good balance (default)
+   - 150+ steps: High quality, slower
+3. **Guidance Scale**:
+   - Lower (3-5): More creative/varied
+   - Default (7): Good balance
+   - Higher (10+): More literal to prompt
+4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
+5. **Seeds**: Use same seed for reproducible results
+
+## Performance
+
+| Inference Steps | Quality | Speed | Use Case |
+|----------------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 (default) | Very Good | Medium | Production |
+| 150+ | Excellent | Slow | Final/critical audio |
+
+## Troubleshooting
+
+### Server not responding
+- Check if server is running: `curl http://localhost:8091/health`
+- Check server logs for errors
+
+### Audio quality issues
+- Increase `num_inference_steps` (e.g., 150)
+- Add negative prompts: `"Low quality, distorted, noisy"`
+- Increase `guidance_scale` for more prompt adherence
+
+### Generation timeout
+- Reduce `num_inference_steps`
+- Reduce `audio_length`
+- Check GPU memory with `nvidia-smi`
+
+### Wrong audio length
+- Ensure `audio_length` is within model limits (~47s max)
+- Adjust `audio_start` if trimming is needed
+
+## See Also
+
+- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
+- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
+- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)

examples/online_serving/stable_audio/curl_examples.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
+
+# Example 1: Simple request with default parameters
+echo "Example 1: Simple request with default parameters"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound audience clapping and cheering in a stadium"
+    }' --output stadium.wav
+
+# Example 2: Request with custom audio_length
+echo "Example 2: Custom audio length (5 seconds)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+
+# Example 3: Request with negative prompt for quality control
+echo "Example 3: With negative prompt"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy"
+    }' --output piano.wav
+
+# Example 4: Full control with all parameters
+echo "Example 4: Full control (custom length, guidance, steps, seed)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+
+# Example 5: Quick generation with fewer steps (faster but lower quality)
+echo "Example 5: Quick generation (fewer steps)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Ocean waves crashing on a beach",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output ocean.wav
+
+echo "All examples completed!"

examples/online_serving/stable_audio/stable_audio_client.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Stable Audio models.
+
+Examples:
+    # Simple generation
+    python stable_audio_client.py --text "The sound of a cat purring"
+
+    # With custom duration
+    python stable_audio_client.py --text "A dog barking" --audio_length 5.0
+
+    # With all parameters
+    python stable_audio_client.py --text "Thunder and rain" \
+        --audio_length 15.0 \
+        --negative_prompt "Low quality" \
+        --guidance_scale 7.0 \
+        --num_inference_steps 100 \
+        --seed 42 \
+        --output thunder.wav
+"""
+
+import argparse
+import sys
+
+import requests
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
+    parser.add_argument(
+        "--api_url",
+        default="http://localhost:8091/v1/audio/generate",
+        help="API endpoint URL",
+    )
+    parser.add_argument(
+        "--text",
+        default="The sound of a cat purring",
+        help="Text prompt for audio generation",
+    )
+    parser.add_argument(
+        "--audio_length",
+        type=float,
+        default=10.0,
+        help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
+    )
+    parser.add_argument(
+        "--audio_start",
+        type=float,
+        default=0.0,
+        help="Audio start time in seconds",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        default="Low quality",
+        help="Negative prompt for classifier-free guidance",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.0,
+        help="Guidance scale for diffusion (higher = more adherence to prompt)",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=100,
+        help="Number of inference steps (higher = better quality, slower)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output",
+        default="stable_audio_output.wav",
+        help="Output file path",
+    )
+    parser.add_argument(
+        "--response_format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm"],
+        help="Audio output format",
+    )
+    return parser.parse_args()
+
+
+def generate_audio(args):
+    """Generate audio using the API."""
+
+    # Build request payload
+    payload = {
+        "input": args.text,
+        "audio_length": args.audio_length,
+        "audio_start": args.audio_start,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.negative_prompt:
+        payload["negative_prompt"] = args.negative_prompt
+    if args.guidance_scale:
+        payload["guidance_scale"] = args.guidance_scale
+    if args.num_inference_steps:
+        payload["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        payload["seed"] = args.seed
+
+    print(f"\n{'=' * 60}")
+    print("Stable Audio - Text-to-Audio Generation")
+    print(f"{'=' * 60}")
+    print(f"API URL: {args.api_url}")
+    print(f"Prompt: {args.text}")
+    print(f"Audio length: {args.audio_length}s")
+    print(f"Negative prompt: {args.negative_prompt}")
+    print(f"Guidance scale: {args.guidance_scale}")
+    print(f"Inference steps: {args.num_inference_steps}")
+    if args.seed is not None:
+        print(f"Seed: {args.seed}")
+    print(f"Output: {args.output}")
+    print(f"{'=' * 60}\n")
+
+    try:
+        # Make the API request
+        print("Generating audio...")
+        response = requests.post(
+            args.api_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=300,  # 5 minute timeout for long generations
+        )
+
+        # Check for errors
+        if response.status_code != 200:
+            print(f"Error: API returned status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return False
+
+        # Save the audio
+        with open(args.output, "wb") as f:
+            f.write(response.content)
+
+        print(f"✓ Audio saved to {args.output}")
+        print(f"  File size: {len(response.content) / 1024:.1f} KB")
+        return True
+
+    except requests.exceptions.Timeout:
+        print("Error: Request timed out. Try reducing inference steps or audio length.")
+        return False
+    except requests.exceptions.ConnectionError:
+        print(f"Error: Could not connect to {args.api_url}")
+        print("Make sure the server is running.")
+        return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def main():
+    args = parse_args()
+    success = generate_audio(args)
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -97,7 +97,6 @@ plugins:
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
-        - "vllm_omni.diffusion.models.dreamid_omni.wan2_2"  # docstring signature warnings break strict docs
         - "vllm_omni.diffusion.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.async_diffusion"  # avoid importing vllm in mkdocs building

pyproject.toml

@@ -32,7 +32,6 @@ classifiers = [
 # Dependencies are now managed dynamically via setup.py based on detected hardware platform.
 # This allows automatic installation of the correct platform-specific dependencies (CUDA/ROCm/CPU/XPU/NPU)
 # without requiring extras like [cuda]. See requirements/ directory for platform-specific dependencies.
-# Note: vllm is intentionally excluded due to entrypoints overwrite issue.
 
 [project.optional-dependencies]
 
@@ -49,7 +48,7 @@ dev = [
     "soundfile>=0.13.1",
     "imageio[ffmpeg]>=0.6.0",
     "opencv-python>=4.12.0.88",
-    "mooncake-transfer-engine==0.3.8.post1",
+    "mooncake-transfer-engine-cuda13>=0.3.9",
     "av", # for ComfyUI tests
     "openpyxl>=3.0.0", # for nightly CI
     "pyttsx3>=2.99",
@@ -90,7 +89,6 @@ Documentation = "https://vllm-omni.readthedocs.io"
 "Bug Tracker" = "https://github.com/vllm-project/vllm-omni/issues"
 
 [project.scripts]
-vllm = "vllm_omni.entrypoints.cli.main:main"
 vllm-omni = "vllm_omni.entrypoints.cli.main:main"
 
 

recipes/Qwen/Qwen3-Omni.md

@@ -1,76 +1,141 @@
-# Qwen3-Omni for multimodal chat on 1x A100 80GB
+# Qwen3-Omni
 
 ## Summary
 
 - Vendor: Qwen
 - Model: `Qwen/Qwen3-Omni-30B-A3B-Instruct`
-- Task: Multimodal chat with text, image, audio, or video input
+- Task: Multimodal chat with text/image/audio/video input
 - Mode: Online serving with the OpenAI-compatible API
 - Maintainer: Community
 
 ## When to use this recipe
 
-Use this recipe when you want a known-good starting point for serving
-`Qwen/Qwen3-Omni-30B-A3B-Instruct` with vLLM-Omni on a single 80 GB A100 and
-validate the deployment with the existing multimodal client examples in this
-repository.
+Use this recipe as a practical baseline for running
+`Qwen/Qwen3-Omni-30B-A3B-Instruct` with the same serving paths already
+exercised by repository examples and tests.
 
 ## References
 
-- Upstream or canonical docs:
+- User guide:
   [`docs/user_guide/examples/online_serving/qwen3_omni.md`](../../docs/user_guide/examples/online_serving/qwen3_omni.md)
-- Related example under `examples/`:
+- Example guide:
   [`examples/online_serving/qwen3_omni/README.md`](../../examples/online_serving/qwen3_omni/README.md)
-- Related issue or discussion:
-  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
 
-## Hardware Support
-
-This recipe currently documents one tested-style reference configuration for
-CUDA GPU serving. Add more sections for other hardware as community validation
-lands.
-
-## GPU
-
-### 1x A100 80GB
-
-#### Environment
+## Environment
 
 - OS: Linux
 - Python: 3.10+
-- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
-- vLLM version: Match the repository requirements for your checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
+- vLLM / vLLM-Omni: use versions from your current checkout, >=0.18.0
 
-#### Command
+## Start server (single command)
 
-Start the server from the repository root:
+From repository root:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
-Async chunking is enabled by default in the bundled deployment config. For
-common runtime tuning, prefer CLI overrides instead of editing or passing a
-custom YAML file:
+Notes:
 
-```bash
-# Disable async chunking for /v1/realtime sessions
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-  --no-async-chunk
-```
+- `--omni` is required.
+- The default deploy config `vllm_omni/deploy/qwen3_omni_moe.yaml` is loaded
+  automatically by model registry.
+- `async_chunk` is enabled by default in this deploy config.
+- Platform deltas under `platforms:` (NPU/ROCm/XPU) are merged automatically on
+  matching runtimes.
 
-Use a custom deploy config only for advanced cases such as custom topology,
-connector wiring, or a larger overlay of stage defaults:
+For advanced customization, pass an overlay YAML:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
   --deploy-config /path/to/your_qwen3_omni_overrides.yaml
 ```
 
+#### Runtime tuning
+
+Prefer CLI overrides for day-to-day tuning:
+
+```bash
+# Disable async chunking when using /v1/realtime
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --no-async-chunk
+
+# Example per-stage tuning in unified launch
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
+
+# Tune max_num_seqs per stage (single process launch)
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{
+    "0": {"max_num_seqs": 8},
+    "1": {"max_num_seqs": 4},
+    "2": {"max_num_seqs": 4}
+  }'
+```
+
+#### Stage-based launch (one stage per process)
+
+Use three terminals (one per stage). Start with the default commands below, then
+add `--max-num-seqs` only if you need explicit per-stage concurrency control.
+
+Default stage-based commands:
+
+```bash
+# Stage 0: Thinker + API server
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --port 8091 \
+  --stage-id 0 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 1: Talker
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 1 \
+  --headless \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 2: Code2Wav
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 2 \
+  --headless \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+```
+
+Optional: explicit per-stage `max_num_seqs` tuning:
+
+```bash
+# Stage 0
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --port 8091 \
+  --stage-id 0 \
+  --max-num-seqs 8 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 1
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 1 \
+  --headless \
+  --max-num-seqs 4 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 2
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 2 \
+  --headless \
+  --max-num-seqs 4 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+```
+
+If you use custom deploy YAML, add `--deploy-config` to each stage command.
+
 #### Verification
 
-Run one of the existing example clients after the server is ready:
+After server startup, run a multimodal example client:
 
 ```bash
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
@@ -80,7 +145,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
   --host localhost
 ```
 
-For a quick API smoke test, request text-only output:
+Quick API smoke test (text-only output):
 
 ```bash
 curl http://localhost:8091/v1/chat/completions \
@@ -92,8 +157,89 @@ curl http://localhost:8091/v1/chat/completions \
   }'
 ```
 
+`{"modalities":["text","audio"]}` means the model returns both text and audio in
+the same response. Use it when you want transcription/content text and TTS audio
+together.
+
+Quick API smoke test (text + audio output):
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+    "messages": [{"role": "user", "content": "Describe vLLM in brief."}],
+    "modalities": ["text", "audio"]
+  }'
+```
+
+Realtime WebSocket check (`/v1/realtime`) requires async chunk disabled:
+
+```bash
+python examples/online_serving/qwen3_omni/openai_realtime_client.py \
+  --url ws://localhost:8091/v1/realtime \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --input-wav /path/to/input_16k_mono.wav \
+  --output-wav realtime_output.wav
+```
+
+#### Benchmark with `vllm bench`
+
+After the server is up, you can run online serving benchmarks with
+`vllm bench serve --omni`.
+
+Text-focused random workload:
+
+```bash
+vllm bench serve \
+  --omni \
+  --host localhost \
+  --port 8091 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --dataset-name random \
+  --num-prompts 40 \
+  --max-concurrency 4 \
+  --random-input-len 2500 \
+  --random-output-len 900 \
+  --ignore-eos \
+  --extra-body '{"modalities":["text"]}' \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+
+If you want benchmark requests to return both text and audio, switch
+`--extra-body` to:
+
+```bash
+--extra-body '{"modalities":["text","audio"]}'
+```
+
+Synthetic multimodal workload (`random-mm`):
+
+```bash
+vllm bench serve \
+  --omni \
+  --host localhost \
+  --port 8091 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --dataset-name random-mm \
+  --num-prompts 40 \
+  --max-concurrency 4 \
+  --random-input-len 2500 \
+  --random-output-len 900 \
+  --random-mm-base-items-per-request 2 \
+  --random-mm-limit-mm-per-prompt '{"image":1,"video":1,"audio":1}' \
+  --random-mm-bucket-config '{"(32, 32, 1)": 0.5, "(0, 1, 1)": 0.5}' \
+  --ignore-eos \
+  --extra-body '{"modalities":["text"]}' \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+
 #### Notes
 
-- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads. Prefer CLI overrides such as `--gpu-memory-utilization` for routine tuning.
-- Key flags: `--omni` is required; async chunking is enabled by default; use `--no-async-chunk` for realtime sessions and `--deploy-config` only for advanced custom deployments.
-- Known limitations: The `/v1/realtime` WebSocket flow is currently unsupported while async chunking is enabled. This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.
+- `/v1/realtime` is unsupported while `async_chunk` is enabled.
+- The default deploy uses `SharedMemoryConnector`; this is for single-host
+  stage wiring.

requirements/cuda.txt

@@ -1,5 +1,5 @@
 -r common.txt
 onnxruntime>=1.23.2
-fa3-fwd==0.0.2
+fa3-fwd==0.0.3
 # nvidia-cublas-cu12 is pinned to 12.9.1.4 via force-reinstall in Dockerfile.ci
 # because torch pins its own cublas version, making it unsatisfiable at resolve time.

requirements/musa.txt

@@ -1,6 +1,6 @@
 -r common.txt
 # MUSA platform dependencies
-torchada>=0.1.50
+torchada>=0.1.52
 onnxruntime>=1.23.2
 mate>=0.2.0
 flash_attn_3>=0.1.4

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -8,8 +8,10 @@ from __future__ import annotations
 
 import importlib.util
 import sys
+import types
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -59,6 +61,7 @@ def mock_tokenizer(mocker):
     tokenizer.all_special_ids = []
     tokenizer.all_special_tokens = []
     tokenizer.vocab_size = 1
+    tokenizer.__len__.return_value = 1
     return tokenizer
 
 
@@ -166,3 +169,58 @@ def test_attach_sets_seed_tts_row_even_without_extra_body():
     row_pos = src.index("seed_tts_row")
     not_ex_pos = src.index("if not ex:")
     assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
+
+
+def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
+    from vllm_omni.benchmarks.data_modules import seed_tts_eval
+
+    calls = {}
+
+    class FakeTensor:
+        def __init__(self, name: str):
+            self.name = name
+            self.device = None
+
+        def to(self, device):
+            self.device = device
+            return self
+
+    class FakeProcessor:
+        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
+            calls["return_attention_mask"] = return_attention_mask
+            assert sampling_rate == 16000
+            assert return_tensors == "pt"
+            assert len(wav) > 0
+            return types.SimpleNamespace(
+                input_features=FakeTensor("features"),
+                attention_mask=FakeTensor("mask") if return_attention_mask else None,
+            )
+
+        def get_decoder_prompt_ids(self, *, language, task):
+            assert language == "english"
+            assert task == "transcribe"
+            return [(1, 2)]
+
+        def batch_decode(self, predicted_ids, *, skip_special_tokens):
+            assert skip_special_tokens
+            assert predicted_ids == [[42]]
+            return ["hello"]
+
+    class FakeModel:
+        def generate(self, input_features, **kwargs):
+            calls["input_features"] = input_features
+            calls["generate_kwargs"] = kwargs
+            return [[42]]
+
+    monkeypatch.setattr(seed_tts_eval, "_ensure_en_asr", lambda: None)
+    monkeypatch.setattr(seed_tts_eval, "_en_processor", FakeProcessor())
+    monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
+    monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
+
+    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
+
+    assert text == "hello"
+    assert calls["return_attention_mask"] is True
+    assert calls["input_features"].device == "cuda:1"
+    assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
+    assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]

tests/config/test_pipeline_registry.py

@@ -6,11 +6,7 @@ from __future__ import annotations
 
 import pytest
 
-from vllm_omni.config.pipeline_registry import (
-    _DIFFUSION_PIPELINES,
-    _OMNI_PIPELINES,
-    _VLLM_OMNI_PIPELINES,
-)
+from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
     PipelineConfig,
@@ -23,17 +19,9 @@ from vllm_omni.config.stage_config import (
 class TestCentralRegistryDeclarations:
     """Every in-tree pipeline must be declared exactly once in the central registry."""
 
-    def test_union_contains_all_omni(self):
+    def test_omni_entries_visible_in_registry(self):
         for key in _OMNI_PIPELINES:
-            assert key in _VLLM_OMNI_PIPELINES
-
-    def test_union_contains_all_diffusion(self):
-        for key in _DIFFUSION_PIPELINES:
-            assert key in _VLLM_OMNI_PIPELINES
-
-    def test_no_duplicate_model_type_between_omni_and_diffusion(self):
-        overlap = set(_OMNI_PIPELINES) & set(_DIFFUSION_PIPELINES)
-        assert not overlap, f"Duplicate model_types across omni/diffusion: {overlap}"
+            assert key in _PIPELINE_REGISTRY
 
     def test_expected_omni_pipelines_present(self):
         # Guard against accidental removal during future refactors.

tests/dfx/conftest.py

@@ -1,13 +1,18 @@
 import json
 import os
+import re
 import subprocess
+import time
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 import pytest
 
+from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
+from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.platforms import current_omni_platform
 
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
@@ -26,18 +31,15 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-def modify_stage(default_path, updates, deletes):
-    kwargs = {}
+def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str:
+    kwargs: dict[str, Any] = {}
     if updates is not None:
         kwargs["updates"] = updates
     if deletes is not None:
         kwargs["deletes"] = deletes
     if kwargs:
-        path = modify_stage_config(default_path, **kwargs)
-    else:
-        path = default_path
-
-    return path
+        return modify_stage_config(default_path, **kwargs)
+    return default_path
 
 
 def _build_serve_args(serve_args: Any) -> list[str]:
@@ -111,6 +113,122 @@ def create_unique_server_params(
     return unique_params
 
 
+def configs_with_platform_stage_configs(configs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Resolve stage config for XPU vs CUDA/ROCm."""
+    out: list[dict[str, Any]] = []
+    for config in configs:
+        config_copy = json.loads(json.dumps(config))
+        if current_omni_platform.is_xpu():
+            config_copy["server_params"]["stage_config_name"] = "xpu/qwen3_omni_ci.yaml"
+        out.append(config_copy)
+    return out
+
+
+def extract_server_args_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[str] | None]:
+    mapping: dict[str, list[str] | None] = {}
+    for cfg in configs:
+        test_name = str(cfg.get("test_name"))
+        server_params = cfg.get("server_params") or {}
+        raw_args = server_params.get("server_args")
+        mapping[test_name] = [str(item) for item in raw_args] if isinstance(raw_args, list) else None
+    return mapping
+
+
+def create_reliability_omni_server_params(
+    configs: list[dict[str, Any]], stage_configs_dir: Path
+) -> list[OmniServerParams]:
+    adjusted_configs = configs_with_platform_stage_configs(configs)
+    unique_params = create_unique_server_params(adjusted_configs, stage_configs_dir)
+    server_args_by_name = extract_server_args_by_test_name(adjusted_configs)
+    return [
+        OmniServerParams(
+            model=model,
+            stage_config_path=stage_config_path,
+            server_args=server_args_by_name.get(test_name),
+        )
+        for test_name, model, stage_config_path, _stage_overrides_json, _extra_cli_args in unique_params
+    ]
+
+
+def supports_video_generation(model_name: str) -> bool:
+    lower = model_name.lower()
+    return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
+
+
+def supports_chat_generation(model_name: str) -> bool:
+    return not supports_video_generation(model_name)
+
+
+def parse_stage_devices(stage_config_path: str) -> str:
+    text = Path(stage_config_path).read_text(encoding="utf-8")
+    raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
+    devices: set[int] = set()
+    for item in raw_devices:
+        for token in item.split(","):
+            token = token.strip()
+            if token:
+                devices.add(int(token))
+    if not devices:
+        raise ValueError(f"No runtime.devices found in stage config: {stage_config_path}")
+    return ",".join(str(x) for x in sorted(devices))
+
+
+def resolve_oom_device_spec(config: dict[str, Any], stage_config_path: str | None) -> str:
+    explicit = config.get("device")
+    if explicit is not None:
+        return str(explicit)
+    if not stage_config_path:
+        return "0"
+    return parse_stage_devices(stage_config_path)
+
+
+def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> None:
+    text = str(exc).lower()
+    assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
+
+
+def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None:
+    """Poll a minimal chat request until success."""
+    deadline = time.time() + timeout_sec
+    payload = json.dumps(
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
+            "stream": False,
+            "modalities": ["text"],
+        }
+    )
+    last_error: str | None = None
+    while time.time() < deadline:
+        try:
+            status, body = post_chat_completions_raw(host, port, payload)
+            if status == 200:
+                return
+            last_error = f"http={status} body={body[:200]!r}"
+        except Exception as exc:  # noqa: BLE001
+            last_error = str(exc)
+        time.sleep(2)
+    raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}")
+
+
+def assert_no_extra_worker_processes(
+    baseline_pids: set[int],
+    worker_pattern: str,
+    timeout_sec: int = 60,
+) -> None:
+    """Ensure no extra worker PIDs remain after teardown."""
+    deadline = time.time() + timeout_sec
+    while time.time() < deadline:
+        current = set(list_remote_process_pids_by_pattern(worker_pattern))
+        extra = current - baseline_pids
+        if not extra:
+            return
+        time.sleep(2)
+    current = set(list_remote_process_pids_by_pattern(worker_pattern))
+    extra = sorted(current - baseline_pids)
+    assert not extra, f"orphan worker processes remain after container teardown: {extra}"
+
+
 def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
     mapping = {}
     for config in configs:

tests/dfx/reliability/README.md

@@ -0,0 +1,14 @@
+# L5(b) Reliability (`tests/dfx/reliability`)
+
+This directory contains reliability fault-injection tests for key models.
+The detailed RFC document is maintained outside this repository (local/internal only).
+
+- Fault helpers: `conftest.py`, `helpers.py`
+- Qwen3 tests: `test_reliability_qwen3_omni.py`
+- Wan2.2 tests: `test_reliability_wan22.py`
+
+```bash
+pytest --collect-only tests/dfx/reliability
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+```

tests/dfx/reliability/conftest.py

@@ -0,0 +1,29 @@
+"""Pytest fixtures for reliability tests."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.dfx.reliability.helpers import FaultInjector
+
+
+@pytest.fixture
+def fault_injector(request: pytest.FixtureRequest) -> FaultInjector:
+    """Indirect only: ``request.param`` must be a ``FaultInjector`` callable."""
+    return request.param
+
+
+@pytest.fixture
+def omni_server_after_fault(omni_server: Any, fault_injector: FaultInjector):
+    """After ``omni_server`` is up, run ``fault_injector(omni_server)``, then yield the server."""
+    fault_injector(omni_server)
+    yield omni_server
+
+
+@pytest.fixture
+def omni_server_after_fault_function(omni_server_function: Any, fault_injector: FaultInjector):
+    """Inject fault after function-scoped server startup, then yield server."""
+    fault_injector(omni_server_function)
+    yield omni_server_function

tests/dfx/reliability/helpers.py

@@ -0,0 +1,567 @@
+"""Shared reliability fault-injection helpers.
+
+This module keeps fault injection callable from tests directly:
+- GPU OOM (CUDA sidecar memory hog)
+- process kill by pattern and signal
+- post-ready hooks via ``fault_injector`` / ``omni_server_after_fault`` fixtures
+"""
+
+from __future__ import annotations
+
+import http.client
+import json
+import logging
+import os
+import select
+import shlex
+import signal
+import subprocess
+import sys
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import psutil
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OomHandle:
+    """Handle for a started CUDA memory hog subprocess."""
+
+    proc: subprocess.Popen | None
+    device: int
+    target_mem_ratio: float
+    start_ts: float
+
+
+def post_chat_completions_raw(
+    host: str,
+    port: int,
+    body: bytes | str,
+    *,
+    content_type: str = "application/json",
+    timeout_sec: int = 120,
+) -> tuple[int, bytes]:
+    """POST /v1/chat/completions with raw bytes; returns (status, response_body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        headers = {"Content-Type": content_type}
+        payload = body.encode("utf-8") if isinstance(body, str) else body
+        conn.request("POST", "/v1/chat/completions", body=payload, headers=headers)
+        resp = conn.getresponse()
+        data = resp.read()
+        return resp.status, data
+    finally:
+        conn.close()
+
+
+def get_health_raw(host: str, port: int, *, timeout_sec: int = 20) -> tuple[int, bytes]:
+    """GET /health with stdlib HTTP client; returns (status, body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        conn.request("GET", "/health")
+        resp = conn.getresponse()
+        return resp.status, resp.read()
+    finally:
+        conn.close()
+
+
+def post_json_raw_http_client(
+    host: str,
+    port: int,
+    path: str,
+    payload: dict[str, Any],
+    *,
+    timeout_sec: int = 30,
+) -> tuple[int, bytes]:
+    """POST JSON to one endpoint with stdlib HTTP client; returns (status, body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        body = json.dumps(payload).encode("utf-8")
+        conn.request("POST", path, body=body, headers={"Content-Type": "application/json"})
+        resp = conn.getresponse()
+        return resp.status, resp.read()
+    finally:
+        conn.close()
+
+
+def post_json_raw(
+    host: str,
+    port: int,
+    path: str,
+    payload: dict[str, Any],
+    *,
+    timeout_sec: int = 30,
+) -> tuple[int, bytes]:
+    """POST JSON to one endpoint; returns (status, body)."""
+    return (
+        post_chat_completions_raw(
+            host,
+            port,
+            json.dumps(payload),
+            content_type="application/json",
+            timeout_sec=timeout_sec,
+        )
+        if path == "/v1/chat/completions"
+        else post_json_raw_http_client(
+            host,
+            port,
+            path,
+            payload,
+            timeout_sec=timeout_sec,
+        )
+    )
+
+
+def extract_openai_error_contract_from_bytes(response_body: bytes) -> dict[str, Any] | None:
+    """Best-effort parse OpenAI-style error object from raw response bytes."""
+    try:
+        payload = json.loads(response_body.decode("utf-8", errors="replace"))
+    except Exception:  # noqa: BLE001
+        return None
+    return extract_openai_error_contract_from_payload(payload)
+
+
+def extract_openai_error_contract_from_payload(payload: Any) -> dict[str, Any] | None:
+    """Best-effort parse OpenAI-style error object from decoded JSON payload."""
+    if not isinstance(payload, dict):
+        return None
+    error_obj = payload.get("error")
+    if not isinstance(error_obj, dict):
+        return None
+    if not isinstance(error_obj.get("message"), str):
+        return None
+    return error_obj
+
+
+def _build_sidecar_cmd(device: int, target_mem_ratio: float, hold_seconds: int, strict: bool) -> list[str]:
+    sidecar = r"""
+import sys
+import time
+import torch
+
+device = int(sys.argv[1])
+target_ratio = float(sys.argv[2])
+hold_seconds = int(sys.argv[3])
+strict = sys.argv[4] == "1"
+
+torch.cuda.init()
+torch.cuda.set_device(device)
+props = torch.cuda.get_device_properties(device)
+free_before, total_bytes = torch.cuda.mem_get_info(device)
+target_bytes = int(free_before * target_ratio)
+chunk_bytes = 256 * 1024 * 1024
+chunks = []
+allocated = 0
+
+while allocated < target_bytes:
+    req_bytes = min(chunk_bytes, target_bytes - allocated)
+    req_elems = max(1, req_bytes // 2)  # float16 -> 2 bytes
+    try:
+        chunk = torch.empty((req_elems,), dtype=torch.float16, device=f"cuda:{device}")
+        chunks.append(chunk)
+        allocated += chunk.numel() * 2
+    except RuntimeError:
+        break
+
+# In strict mode, keep filling with smaller chunks until allocator rejects.
+# This minimizes residual free memory and makes fault-path assertions steadier.
+if strict:
+    tail_chunk_bytes = [64 * 1024 * 1024, 16 * 1024 * 1024, 4 * 1024 * 1024, 1 * 1024 * 1024]
+    for tail_bytes in tail_chunk_bytes:
+        while True:
+            req_elems = max(1, tail_bytes // 2)
+            try:
+                chunk = torch.empty((req_elems,), dtype=torch.float16, device=f"cuda:{device}")
+                chunks.append(chunk)
+                allocated += chunk.numel() * 2
+            except RuntimeError:
+                break
+
+achieved_ratio = allocated / max(1, props.total_memory)
+achieved_free_ratio = allocated / max(1, int(free_before))
+free_after, _ = torch.cuda.mem_get_info(device)
+if strict and allocated < target_bytes:
+    print(
+        "ERROR:"
+        f"achieved_free_ratio={achieved_free_ratio:.4f};"
+        f"achieved_total_ratio={achieved_ratio:.4f};"
+        f"free_before={int(free_before)};"
+        f"free_after={int(free_after)};"
+        f"target_bytes={target_bytes};"
+        f"allocated={allocated}",
+        flush=True,
+    )
+    sys.exit(2)
+
+print(
+    "READY:"
+    f"achieved_free_ratio={achieved_free_ratio:.4f};"
+    f"achieved_total_ratio={achieved_ratio:.4f};"
+    f"free_before={int(free_before)};"
+    f"free_after={int(free_after)};"
+    f"target_bytes={target_bytes};"
+    f"allocated={allocated}",
+    flush=True,
+)
+if hold_seconds <= 0:
+    while True:
+        time.sleep(3600)
+time.sleep(hold_seconds)
+print("DONE", flush=True)
+"""
+    return [
+        sys.executable,
+        "-u",
+        "-c",
+        sidecar,
+        str(device),
+        str(target_mem_ratio),
+        str(hold_seconds),
+        "1" if strict else "0",
+    ]
+
+
+def start_gpu_oom_hog(
+    *,
+    device: int = 0,
+    target_mem_ratio: float = 0.95,
+    hold_seconds: int = 60,
+    startup_timeout_sec: int = 20,
+    strict: bool = True,
+    poll_interval_sec: float = 0.2,
+) -> OomHandle:
+    """Start a CUDA sidecar process that occupies GPU memory to trigger OOM.
+
+    Note:
+        ``target_mem_ratio`` is evaluated against free memory at injection start
+        (not total GPU memory), i.e. success gate is ``allocated / free_before``.
+        ``hold_seconds <= 0`` means keeping OOM pressure until the sidecar is
+        explicitly stopped via ``stop_gpu_oom_hog(s)``.
+    """
+    if os.name == "nt":
+        raise RuntimeError("CUDA OOM sidecar is intended for Linux CI/runtime.")
+    if not (0.0 <= target_mem_ratio < 1.0):
+        raise ValueError("target_mem_ratio should be in [0.0, 1.0).")
+
+    # Explicit opt-out for debugging: keep API shape stable while disabling injection.
+    if target_mem_ratio == 0.0:
+        print(f"[oom-sidecar][gpu={device}] DISABLED: target_mem_ratio=0.0 (no OOM injection)", flush=True)
+        return OomHandle(
+            proc=None,
+            device=device,
+            target_mem_ratio=target_mem_ratio,
+            start_ts=time.time(),
+        )
+
+    cmd = _build_sidecar_cmd(device, target_mem_ratio, hold_seconds, strict)
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+    assert proc.stdout is not None
+
+    deadline = time.time() + startup_timeout_sec
+    logs: list[str] = []
+    while time.time() < deadline:
+        ready, _, _ = select.select([proc.stdout], [], [], poll_interval_sec)
+        if ready:
+            line = proc.stdout.readline().strip()
+            if line:
+                logs.append(line)
+                print(f"[oom-sidecar][gpu={device}] {line}", flush=True)
+                if line.startswith("READY:"):
+                    return OomHandle(
+                        proc=proc,
+                        device=device,
+                        target_mem_ratio=target_mem_ratio,
+                        start_ts=time.time(),
+                    )
+                if line.startswith("ERROR:"):
+                    proc.terminate()
+                    raise RuntimeError(f"OOM sidecar failed to reach target: {line}")
+        if proc.poll() is not None:
+            break
+
+    proc.terminate()
+    if logs:
+        print(f"[oom-sidecar][gpu={device}] startup logs: {' | '.join(logs)}", flush=True)
+    raise TimeoutError(f"OOM sidecar startup timeout. logs={logs}")
+
+
+def stop_gpu_oom_hog(handle: OomHandle, *, timeout_sec: int = 5) -> None:
+    """Stop and cleanup CUDA OOM sidecar."""
+    proc = handle.proc
+    if proc is None:
+        return
+    if proc.poll() is not None:
+        return
+    proc.terminate()
+    try:
+        proc.wait(timeout=timeout_sec)
+    except subprocess.TimeoutExpired:
+        proc.kill()
+        proc.wait(timeout=timeout_sec)
+
+
+def inject_gpu_oom(
+    *,
+    device: int | str | list[int] = 0,
+    target_mem_ratio: float = 0.95,
+    hold_seconds: int = 60,
+    startup_timeout_sec: int = 20,
+    strict: bool = True,
+) -> OomHandle | list[OomHandle]:
+    """Convenience wrapper to start CUDA OOM sidecar(s).
+
+    Args:
+        device: One device id (``0``), comma-separated string (``"0,1,2"``),
+            or a list of device ids (``[0, 1, 2]``).
+        hold_seconds: OOM hold time in seconds; ``<=0`` keeps pressure until
+            ``stop_gpu_oom_hogs`` is called.
+    """
+    if isinstance(device, int):
+        devices = [device]
+    elif isinstance(device, str):
+        devices = [int(x.strip()) for x in device.split(",") if x.strip()]
+    else:
+        devices = [int(x) for x in device]
+    if not devices:
+        raise ValueError("device must not be empty.")
+
+    handles = [
+        start_gpu_oom_hog(
+            device=dev,
+            target_mem_ratio=target_mem_ratio,
+            hold_seconds=hold_seconds,
+            startup_timeout_sec=startup_timeout_sec,
+            strict=strict,
+        )
+        for dev in devices
+    ]
+    if len(handles) == 1:
+        return handles[0]
+    return handles
+
+
+def stop_gpu_oom_hogs(handles: OomHandle | list[OomHandle], *, timeout_sec: int = 5) -> None:
+    """Stop one or multiple OOM sidecars."""
+    if isinstance(handles, OomHandle):
+        stop_gpu_oom_hog(handles, timeout_sec=timeout_sec)
+        return
+    for handle in handles:
+        stop_gpu_oom_hog(handle, timeout_sec=timeout_sec)
+
+
+def _runtime_teardown_ssh_target() -> str:
+    target = os.getenv("RUNTIME_TEARDOWN_SSH_TARGET", "").strip()
+    # Default to root@127.0.0.1 for same-host SSH control path.
+    return target or "root@127.0.0.1"
+
+
+def _runtime_teardown_ssh_cmd(remote_cmd: str, *, step: str | None = None) -> subprocess.CompletedProcess[str]:
+    ssh_target = _runtime_teardown_ssh_target()
+    default_reuse_opts = "-o ControlMaster=auto -o ControlPersist=10m -o ControlPath=/tmp/vllm-rt-ssh-%r@%h:%p"
+    raw_opts = os.getenv("RUNTIME_TEARDOWN_SSH_OPTS", "").strip()
+    ssh_opts = shlex.split(raw_opts or default_reuse_opts)
+    timeout_sec = int(os.getenv("RUNTIME_TEARDOWN_SSH_TIMEOUT_SEC", "600"))
+    step_prefix = f"[runtime-teardown][ssh]{f'[{step}]' if step else ''}"
+    print(f"{step_prefix} target={ssh_target} running remote command...", flush=True)
+    # IMPORTANT: SSH joins remote argv into one shell command string. If we pass
+    # ["bash", "-c", remote_cmd] as separate argv items, remote shell parsing can
+    # make `-c` consume only the first word (e.g. "docker"), causing docker help.
+    # Wrap the whole command as one quoted string for remote bash -c.
+    remote_invocation = f"bash --noprofile --norc -c {shlex.quote(remote_cmd)}"
+    try:
+        out = subprocess.run(
+            ["ssh", *ssh_opts, ssh_target, remote_invocation],
+            check=False,
+            capture_output=True,
+            text=True,
+            timeout=timeout_sec,
+        )
+    except subprocess.TimeoutExpired as exc:
+        raise RuntimeError(
+            f"{step_prefix} timed out after {timeout_sec}s. Increase RUNTIME_TEARDOWN_SSH_TIMEOUT_SEC if needed."
+        ) from exc
+    print(f"{step_prefix} exit_code={out.returncode}", flush=True)
+    return out
+
+
+def list_remote_process_pids_by_pattern(pattern: str) -> list[int]:
+    """Return matched PIDs from remote host ``pgrep -f <pattern>`` via SSH."""
+    cmd = f"pgrep -f {shlex.quote(pattern)} || true"
+    out = _runtime_teardown_ssh_cmd(cmd, step="pgrep")
+    if out.returncode not in (0, 1):
+        raise RuntimeError(f"remote pgrep failed for pattern={pattern!r}: {out.stderr.strip()}")
+    return [int(item) for item in out.stdout.split() if item.strip().isdigit()]
+
+
+def inject_process_kill(
+    *,
+    grep_pattern: str,
+    signal_name: str = "SIGTERM",
+    limit: int | None = None,
+    allow_zero_match: bool = False,
+    execute_kill: bool = True,
+) -> list[int]:
+    """Kill processes matching pattern with selected signal."""
+    if os.name == "nt":
+        raise RuntimeError("process-kill helper currently supports POSIX platforms only.")
+    if not grep_pattern.strip():
+        raise ValueError("grep_pattern must not be empty.")
+
+    sig = getattr(signal, signal_name, None)
+    if sig is None:
+        raise ValueError(f"Unsupported signal_name: {signal_name}")
+
+    out = subprocess.run(
+        ["pgrep", "-f", grep_pattern],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    pids = [int(item) for item in out.stdout.split() if item.strip().isdigit()]
+    if limit is not None:
+        pids = pids[:limit]
+
+    if not pids and not allow_zero_match:
+        raise RuntimeError(f"No process matched pattern: {grep_pattern}")
+
+    if execute_kill:
+        for pid in pids:
+            os.kill(pid, sig)
+    return pids
+
+
+def _safe_proc_info(pid: int) -> tuple[str, str]:
+    """Best-effort process name/cmdline lookup for debug logging."""
+    try:
+        proc = psutil.Process(pid)
+        name = proc.name()
+        cmdline = " ".join(proc.cmdline()) or "<empty-cmdline>"
+        return name, cmdline
+    except Exception:  # noqa: BLE001
+        return "<unknown>", "<unavailable>"
+
+
+def _list_server_process_tree(server: Any) -> list[int]:
+    """Return [root, descendants...] PIDs for the current test server instance."""
+    root_proc = getattr(server, "proc", None)
+    if root_proc is None or getattr(root_proc, "pid", None) is None:
+        return []
+
+    root_pid = int(root_proc.pid)
+    try:
+        root = psutil.Process(root_pid)
+    except Exception:  # noqa: BLE001
+        return [root_pid]
+
+    descendants = [child.pid for child in root.children(recursive=True)]
+    return [root_pid, *descendants]
+
+
+def _log_server_process_tree(server: Any) -> None:
+    """Print server process tree for debugging fault injection targets."""
+    pids = _list_server_process_tree(server)
+    if not pids:
+        logger.warning("[reliability][process-kill] current server has no visible process tree")
+        return
+    for pid in pids:
+        name, cmdline = _safe_proc_info(pid)
+        print(
+            f"[reliability][process-kill] current_server_proc pid={pid} name={name} cmdline={cmdline}",
+            flush=True,
+        )
+
+
+FaultInjector = Callable[[Any], None]
+"""Callable invoked with the live ``OmniServer`` after it is ready (see ``omni_server_after_fault``)."""
+
+
+def make_process_kill_fault_injector(
+    *,
+    grep_patterns: str | Sequence[str],
+    signal_name: str = "SIGKILL",
+    limit: int = 1,
+    post_kill_wait_seconds: float = 0.0,
+) -> FaultInjector:
+    """Build a post-ready injector that kills processes matched by ``pgrep -f``.
+
+    Tries each pattern in order until at least one PID is killed. If none match,
+    the returned callable issues ``pytest.skip`` (same behavior as the previous
+    inline reliability test).
+
+    Args:
+        grep_patterns: One pattern or an ordered list of patterns.
+        signal_name: Passed to :func:`inject_process_kill` (e.g. ``SIGKILL``).
+        limit: Maximum PIDs to kill per pattern (default ``1``).
+        post_kill_wait_seconds: Optional wait time after kill before test request starts.
+    """
+    patterns: tuple[str, ...] = (grep_patterns,) if isinstance(grep_patterns, str) else tuple(grep_patterns)
+
+    def _inject(server: Any) -> None:
+        _log_server_process_tree(server)
+        server_tree = set(_list_server_process_tree(server))
+        if not server_tree:
+            logger.warning(
+                "[reliability][process-kill] no server process tree found; fallback to global pgrep matching"
+            )
+        for pattern in patterns:
+            print(
+                f"[reliability][process-kill] trying pattern={pattern} signal={signal_name} limit={limit}",
+                flush=True,
+            )
+            pids = inject_process_kill(
+                grep_pattern=pattern,
+                signal_name=signal_name,
+                limit=limit,
+                allow_zero_match=True,
+                execute_kill=False,
+            )
+            filtered = [pid for pid in pids if not server_tree or pid in server_tree]
+            if pids and not filtered:
+                logger.warning(
+                    "[reliability][process-kill] pattern=%s matched non-server pids=%s, skip them",
+                    pattern,
+                    pids,
+                )
+                continue
+            if filtered:
+                sig = getattr(signal, signal_name, None)
+                if sig is None:
+                    raise ValueError(f"Unsupported signal_name: {signal_name}")
+                for pid in filtered:
+                    name, cmdline = _safe_proc_info(pid)
+                    print(
+                        f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
+                        flush=True,
+                    )
+                    os.kill(pid, sig)
+                print(
+                    f"[reliability][process-kill] matched pattern={pattern} killed_pids={filtered} killed_count={len(filtered)}",
+                    flush=True,
+                )
+                if post_kill_wait_seconds > 0:
+                    print(
+                        f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
+                        flush=True,
+                    )
+                    time.sleep(post_kill_wait_seconds)
+                return
+        logger.warning(
+            "[reliability][process-kill] no process matched patterns=%s signal=%s limit=%s",
+            patterns,
+            signal_name,
+            limit,
+        )
+        pytest.skip("no matching runtime process found for kill injection")
+
+    return _inject

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -0,0 +1,682 @@
+"""
+Qwen3-Omni reliability integration tests.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import errno
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Protocol
+
+import pytest
+import torch
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+)
+from tests.dfx.reliability.helpers import (
+    extract_openai_error_contract_from_bytes,
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    post_chat_completions_raw,
+    post_json_raw,
+    stop_gpu_oom_hogs,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import dummy_messages_from_mix_data
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "qwen3_omni_reliability_async_chunk",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_config_name": "qwen3_omni_moe.yaml",
+            "server_args": ["--async-chunk"],
+        },
+    },
+    {
+        "test_name": "qwen3_omni_reliability_default",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_config_name": "qwen3_omni_moe.yaml",
+            "server_args": ["--no-async-chunk"],
+        },
+    },
+]
+
+DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
+
+
+def _default_oom_device_spec() -> str:
+    """Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
+    count = torch.cuda.device_count()
+    if count <= 0:
+        return "0"
+    return ",".join(str(i) for i in range(count))
+
+
+OOM_INJECTION_CONFIG = {
+    "device": _default_oom_device_spec(),
+    "target_mem_ratio": 0.92,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": False,
+}
+# Post-fault recovery probe: keep the API process alive on the same GPU.
+# ``strict=True`` sidecars keep allocating until CUDA refuses, which often
+# evicts/kills the server; here we only take a fraction of *initial* free
+# memory and stop, then rely on a heavy multimodal forward to hit OOM.
+OOM_RECOVER_INJECTION_CONFIG = {
+    "device": _default_oom_device_spec(),
+    # ``strict=False`` stops at this fraction of *initial* free memory per GPU (see helpers sidecar).
+    # On multi-GPU hosts a mild ratio plus a small e2e-style mix can still leave several GB free
+    # on a data-parallel device and the request succeeds; 0.92+ pairs better with the heavy
+    # probe payload below. Raise toward 0.94 only if fault phase still never fails; avoid strict=True
+    # here so the API process is not squeezed to zero slack.
+    "target_mem_ratio": 0.8,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": False,
+}
+FAULT_ERROR_KEYWORDS = (
+    "the request failed",
+    "oom",
+    "out of memory",
+    "cuda",
+    "orchestrator",
+    "timeout",
+    "connection",
+    "500",
+    "503",
+)
+RUNTIME_WORKER_PATTERN = "VLLM::Worker"
+
+
+class _HasServeArgs(Protocol):
+    model: str
+    serve_args: list[str]
+
+
+def _get_system_prompt() -> dict:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _get_mix_prompt() -> str:
+    return "What is recited in the audio? What is in this image? Describe the video briefly."
+
+
+def _mix_chat_completions_probe_payload(omni_server: _HasServeArgs) -> dict[str, Any]:
+    """Heavy multimodal ``/v1/chat/completions`` body for OOM recover tests.
+
+    Uses the same synthetic sizes as ``test_reliability_fault_gpu_oom_chat_large_payload_failure``
+    (large vision tensors + long text), not the light expansion mix (224² / 5s audio): under
+    ``strict=False`` sidecars, a small mix often still completes on a GPU that retains multi-GB
+    free after the hog stops at ``target_mem_ratio``.
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text="What is recited in the audio? What is in this image? What is in this video? " * 200,
+    )
+    return {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text", "audio"],
+    }
+
+
+def _fault_keywords_match_response(*, status: int, body: bytes) -> bool:
+    """True when HTTP status or body text looks like an expected fault / OOM class outcome."""
+    if status >= 400:
+        return True
+    text = body.decode("utf-8", errors="replace").lower()
+    # Avoid matching bare ``"500"`` / ``"503"`` in successful streamed bodies (base64/SSE noise).
+    body_fault_hints = (
+        "the request failed",
+        "oom",
+        "out of memory",
+        "cuda",
+        "orchestrator",
+        "timeout",
+        "connection refused",
+        "connection reset",
+    )
+    return any(key in text for key in body_fault_hints)
+
+
+def _stage_config_path_from_omni_server(omni_server: _HasServeArgs) -> str | None:
+    args: list[str] = omni_server.serve_args
+    for i, arg in enumerate(args):
+        if arg == "--stage-configs-path" and i + 1 < len(args):
+            return args[i + 1]
+        if arg.startswith("--stage-configs-path="):
+            return arg.split("=", 1)[1]
+    return None
+
+
+def _looks_like_server_unreachable(exc: BaseException) -> bool:
+    """True when /health cannot be reached because nothing is listening (process exited)."""
+    if isinstance(exc, (ConnectionRefusedError, BrokenPipeError, ConnectionResetError)):
+        return True
+    errno_val = getattr(exc, "errno", None)
+    if isinstance(exc, OSError) and errno_val is not None:
+        return errno_val in (
+            errno.ECONNREFUSED,
+            errno.ECONNRESET,
+            errno.EPIPE,
+        )
+    msg = str(exc).lower()
+    return "connection refused" in msg or "actively refused" in msg
+
+
+QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
+    omni_server_function,
+) -> None:
+    """Black-box: text chat vs omni audio output should expose a consistent error contract under OOM.
+
+    Speech-style output is requested via ``/v1/chat/completions`` with ``modalities`` that include ``audio``.
+    """
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    mix_base = _mix_chat_completions_probe_payload(omni_server_function)
+    oom_contract_timeout_sec = 120
+    try:
+        chat_status, chat_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            {
+                "model": mix_base["model"],
+                "messages": mix_base["messages"],
+                "stream": False,
+                "modalities": ["text"],
+            },
+            timeout_sec=oom_contract_timeout_sec,
+        )
+        speech_status, speech_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            {
+                "model": mix_base["model"],
+                "messages": mix_base["messages"],
+                "stream": False,
+                "modalities": ["text", "audio"],
+            },
+            timeout_sec=oom_contract_timeout_sec,
+        )
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+    # Under OOM pressure both chat and speech should surface runtime-class (5xx)
+    # failures instead of request-validation-class (4xx) errors.
+    chat_error = extract_openai_error_contract_from_bytes(chat_body)
+    speech_error = extract_openai_error_contract_from_bytes(speech_body)
+    print(chat_status, chat_error, speech_status, speech_error)
+
+    assert chat_status >= 500, f"expected chat error under OOM, got status={chat_status}"
+    assert speech_status >= 500, f"expected speech runtime error under OOM, got status={speech_status}"
+
+    assert chat_error is not None, f"chat error payload not OpenAI-compatible: {chat_body[:300]!r}"
+    assert speech_error is not None, f"speech error payload not OpenAI-compatible: {speech_body[:300]!r}"
+    assert "code" in chat_error, f"chat error lacks code field: {chat_error!r}"
+    assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+        messages = dummy_messages_from_mix_data(
+            system_prompt=_get_system_prompt(),
+            video_data_url=video_data_url,
+            image_data_url=image_data_url,
+            audio_data_url=audio_data_url,
+            content_text=f"{_get_mix_prompt()} " * 200,
+        )
+        request_config = {
+            "model": omni_server_function.model,
+            "messages": messages,
+            "stream": True,
+            "key_words": {"audio": ["test"]},
+        }
+        try:
+            openai_client_function.send_omni_request(request_config, request_num=1)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected large chat payload request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+        messages = dummy_messages_from_mix_data(
+            system_prompt=_get_system_prompt(),
+            video_data_url=video_data_url,
+            image_data_url=image_data_url,
+            audio_data_url=audio_data_url,
+            content_text=f"{_get_mix_prompt()} " * 200,
+        )
+        request_config = {
+            "model": omni_server_function.model,
+            "messages": messages,
+            "stream": True,
+            "modalities": ["text", "audio"],
+            "key_words": {"audio": ["test"]},
+        }
+        try:
+            openai_client_function.send_omni_request(request_config, request_num=4)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected concurrent request failure under OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_request_failure(
+    omni_server_after_fault_function, openai_client_function
+) -> None:
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_get_system_prompt(),
+        content_text="What is the capital of China? Answer in 20 words.",
+    )
+    request_config = {
+        "model": omni_server_after_fault_function.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+    try:
+        openai_client_function.send_omni_request(request_config, request_num=1)
+    except Exception as exc:
+        assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+) -> None:
+    """Black-box: after worker SIGKILL, /health→503, chat fails fast, concurrent chat does not hang."""
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    model = omni_server_after_fault_function.model
+
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    saw_503 = False
+    health_final_status: int | None = None
+    health_final_body = b""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            health_final_status, health_final_body = status, body
+            if status == 503:
+                saw_503 = True
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    assert saw_503, (
+        f"[process_kill health] expected /health to become 503 after fault injection, got {last_observation}"
+    )
+
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
+        "stream": False,
+        "modalities": ["text"],
+    }
+    ff_status: int | None = None
+    ff_body = b""
+    ff_exc: BaseException | None = None
+    start = time.monotonic()
+    try:
+        ff_status, ff_body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        assert ff_status >= 500, (
+            f"[process_kill fast_fail] expected server-side failure after fault, "
+            f"got status={ff_status}, body={ff_body[:200]!r}"
+        )
+    except Exception as exc:
+        ff_exc = exc
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
+
+    payload_json = json.dumps(
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": "What is the capital of China? Answer in one word."}],
+            "stream": False,
+            "modalities": ["text"],
+        }
+    )
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        futures = [
+            executor.submit(
+                post_chat_completions_raw,
+                host,
+                port,
+                payload_json,
+                timeout_sec=20,
+            )
+            for _ in range(4)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=30,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time requests hung: pending={len(pending)}"
+    assert elapsed < 30, f"[process_kill concurrent] fault-time request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    conc_debug: list[Any] = []
+    for future in done:
+        try:
+            status, body = future.result()
+            conc_debug.append((status, body[:200]))
+            if status >= 500:
+                fault_observed = True
+        except Exception as exc:
+            conc_debug.append(repr(exc))
+            fault_observed = True
+    # DEBUG: remove before merge
+    print(
+        health_final_status,
+        health_final_body[:200],
+        ff_status,
+        ff_body[:200],
+        ff_exc,
+        conc_debug,
+    )
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one request to fail after process-kill fault injection"
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="issue#2327")
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
+    omni_server_function,
+) -> None:
+    """Black-box: under bounded GPU pressure, a heavy mix chat fails; after hog stops, the same request succeeds.
+
+    Uses a milder sidecar than ``OOM_INJECTION_CONFIG`` (``strict=False`` and a capped
+    ``target_mem_ratio``) so the serving process stays alive: pressure should starve
+    the multimodal forward pass, not exhaust the whole device allocator.
+
+    Fault and recovery phases use raw ``/v1/chat/completions`` HTTP so tests can assert
+    on status codes and JSON error payloads (or fault keywords in streamed bodies).
+    """
+    device_spec = resolve_oom_device_spec(
+        OOM_RECOVER_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_RECOVER_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_RECOVER_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_RECOVER_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_RECOVER_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    mix_payload = _mix_chat_completions_probe_payload(omni_server_function)
+    mix_chat_timeout_sec = 180
+    print(
+        "[oom-recover] injection ready "
+        f"device={device_spec} target_mem_ratio={OOM_RECOVER_INJECTION_CONFIG['target_mem_ratio']} "
+        f"strict={OOM_RECOVER_INJECTION_CONFIG['strict']} hold_seconds={OOM_RECOVER_INJECTION_CONFIG['hold_seconds']}"
+    )
+    print(f"[oom-recover] target endpoint http://{host}:{port}/v1/chat/completions")
+
+    failure_observed = False
+    fault_status: int | None = None
+    fault_body = b""
+    try:
+        try:
+            print(f"[oom-recover] fault-phase request start timeout={mix_chat_timeout_sec}s")
+            fault_status, fault_body = post_json_raw(
+                host,
+                port,
+                "/v1/chat/completions",
+                mix_payload,
+                timeout_sec=mix_chat_timeout_sec,
+            )
+            print(f"[oom-recover] fault-phase request done status={fault_status} body_prefix={fault_body[:200]!r}")
+        except Exception as exc:
+            failure_observed = True
+            print(f"[oom-recover] fault-phase request raised {type(exc).__name__}: {exc!r}")
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            assert fault_status is not None
+            failure_observed = fault_status != 200 or _fault_keywords_match_response(
+                status=fault_status,
+                body=fault_body,
+            )
+            assert failure_observed, (
+                "expected mix multimodal request failure while OOM pressure is active; "
+                f"status={fault_status}, body_prefix={fault_body[:500]!r}"
+            )
+            if fault_status >= 400:
+                err = extract_openai_error_contract_from_bytes(fault_body)
+                if err is None:
+                    assert _fault_keywords_match_response(status=fault_status, body=fault_body), (
+                        "non-2xx without OpenAI-style error object should still mention fault hints in body; "
+                        f"status={fault_status}, body_prefix={fault_body[:500]!r}"
+                    )
+                else:
+                    assert isinstance(err.get("message"), str) and err["message"].strip(), (
+                        f"structured error should carry a non-empty message: {err!r}"
+                    )
+            else:
+                assert _fault_keywords_match_response(status=fault_status, body=fault_body), (
+                    "HTTP 200 under pressure should still surface fault hints in the response body; "
+                    f"body_prefix={fault_body[:500]!r}"
+                )
+    finally:
+        print("[oom-recover] stopping OOM sidecar(s)")
+        stop_gpu_oom_hogs(handle)
+        print("[oom-recover] OOM sidecar(s) stopped")
+
+    recovery_deadline = time.monotonic() + 90.0
+    terminal_health: int | None = None
+    unreachable_streak = 0
+    last_health_exc: BaseException | None = None
+    while time.monotonic() < recovery_deadline:
+        try:
+            status, _ = get_health_raw(host, port, timeout_sec=5)
+            unreachable_streak = 0
+            if status in (200, 503):
+                terminal_health = status
+                print(f"[oom-recover] terminal health reached status={terminal_health}")
+                break
+        except Exception as exc:
+            last_health_exc = exc
+            if _looks_like_server_unreachable(exc):
+                unreachable_streak += 1
+                if unreachable_streak >= 5:
+                    pytest.fail(
+                        "after OOM sidecar stopped, /health is unreachable (connection refused / reset). "
+                        "The APIServer process likely exited (e.g. orchestrator thread crash under OOM); "
+                        "this test expects the server to stay up for post-fault health polling. "
+                        f"last_exc={last_health_exc!r}"
+                    )
+            else:
+                unreachable_streak = 0
+        time.sleep(1.0)
+    else:
+        pytest.fail(
+            "server did not converge to a terminal health state after OOM pressure was removed; "
+            f"last_health_exc={last_health_exc!r}"
+        )
+
+    probe_payload = {
+        "model": omni_server_function.model,
+        "messages": [{"role": "user", "content": "What is the capital of China? Answer in one word."}],
+        "stream": False,
+        "modalities": ["text"],
+    }
+    start = time.monotonic()
+    post_mix_status: int | None = None
+    post_mix_body = b""
+    request_status: int | None = None
+    try:
+        assert terminal_health is not None
+        if terminal_health == 200:
+            print(f"[oom-recover] recovery-phase mix request start timeout={mix_chat_timeout_sec}s")
+            post_mix_status, post_mix_body = post_json_raw(
+                host,
+                port,
+                "/v1/chat/completions",
+                mix_payload,
+                timeout_sec=mix_chat_timeout_sec,
+            )
+            print(
+                "[oom-recover] recovery-phase mix request done "
+                f"status={post_mix_status} body_prefix={post_mix_body[:200]!r}"
+            )
+        else:
+            print("[oom-recover] terminal health=503, send lightweight probe request")
+            request_status, _ = post_json_raw(host, port, "/v1/chat/completions", probe_payload, timeout_sec=20)
+            print(f"[oom-recover] lightweight probe done status={request_status}")
+    except Exception:
+        post_mix_status = None
+        request_status = None
+    elapsed = time.monotonic() - start
+    assert elapsed < 200, f"post-fault request should not hang after OOM removal: {elapsed:.2f}s"
+
+    if terminal_health == 200:
+        assert post_mix_status == 200, (
+            "health recovered but repeated mix multimodal request did not return HTTP 200; "
+            f"status={post_mix_status}, body_prefix={post_mix_body[:600]!r}"
+        )
+        recover_err = extract_openai_error_contract_from_bytes(post_mix_body)
+        assert recover_err is None, f"unexpected error object after recovery: {recover_err!r}"
+    else:
+        assert request_status is None or request_status >= 500, (
+            "unhealthy terminal state should fail fast on requests, "
+            f"got health={terminal_health}, request_status={request_status}"
+        )

tests/dfx/reliability/test_reliability_wan22.py

@@ -0,0 +1,411 @@
+"""
+Wan2.2 reliability integration tests.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+    supports_video_generation,
+)
+from tests.dfx.reliability.helpers import (
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    stop_gpu_oom_hogs,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "wan22_t2v_reliability_default",
+        "server_params": {
+            "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "server_args": [
+                "--num-gpus",
+                "1",
+                "--boundary-ratio",
+                "0.875",
+                "--flow-shift",
+                "5.0",
+                "--disable-log-stats",
+            ],
+        },
+    }
+]
+
+E2E_STAGE_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "e2e" / "stage_configs"
+OOM_INJECTION_CONFIG = {
+    "target_mem_ratio": 0.95,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": True,
+}
+FAULT_ERROR_KEYWORDS = (
+    "oom",
+    "out of memory",
+    "cuda",
+    "job failed",
+    "unknown error",
+    "internal server error",
+    "500 server error",
+)
+PROCESS_KILL_ERROR_KEYWORDS = (
+    "timeout",
+    "did not complete within",
+    "connection",
+    "engine",
+    "orchestrator",
+    "dead",
+    "internal",
+    "500",
+    "503",
+)
+
+WAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, E2E_STAGE_CONFIGS_DIR)
+DIFFUSION_VIDEO_PARAMS = [param for param in WAN_PARAMS if supports_video_generation(param.model)]
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, openai_client_function) -> None:
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        request_config = {
+            "form_data": {
+                "prompt": "Generate a realistic road-driving video with camera motion.",
+                "width": 512,
+                "height": 512,
+                "fps": 8,
+                "num_frames": 8,
+                "guidance_scale": 1.0,
+                "flow_shift": 5.0,
+                "num_inference_steps": 4,
+                "seed": 42,
+            },
+            "image_reference": image_data_url,
+            "stream": False,
+        }
+        try:
+            openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected large /v1/videos request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_video_request_failure(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    request_config = {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
+    try:
+        openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+    except Exception as exc:
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected /v1/videos request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    """Black-box: after process kill, /v1/videos fails fast and concurrent calls don't hang.
+
+    /health→503 is checked last: optional ``pytest.skip`` then ``assert`` (drop the ``if`` to harden).
+    """
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    url = f"http://{host}:{port}/v1/videos"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    request_config = {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
+
+    # Phase-1: one new /v1/videos request should fail fast.
+    payload = {
+        "prompt": "fast-fail check",
+        "width": "512",
+        "height": "512",
+        "fps": "8",
+        "num_frames": "8",
+        "num_inference_steps": "4",
+    }
+    start = time.monotonic()
+    try:
+        response = requests.post(
+            url,
+            data=payload,
+            headers={"Accept": "application/json"},
+            timeout=20,
+        )
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] /v1/videos did not fail fast after fault: {elapsed:.2f}s"
+        assert response.status_code >= 500, (
+            "[process_kill fast_fail] expected server-side error after fatal fault, "
+            f"got status={response.status_code}, body={response.text[:200]!r}"
+        )
+    except Exception:
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] /v1/videos exception was too slow after fault: {elapsed:.2f}s"
+
+    # Phase-2: concurrent requests should converge and at least one should fail.
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(openai_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=40,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time video requests hung: pending={len(pending)}"
+    assert elapsed < 40, f"[process_kill concurrent] fault-time video request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    conc_debug: list[Any] = []
+    for future in done:
+        try:
+            future.result()
+            conc_debug.append("ok")
+        except Exception as exc:
+            fault_observed = True
+            conc_debug.append(repr(exc))
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one /v1/videos request to fail after fault; "
+        f"conc_debug={conc_debug}"
+    )
+
+    # Phase-3: /health→503 after fault. Optional skip today; remove the ``if`` later to harden.
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    final_health_status: int | None = None
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            final_health_status = status
+            if status == 503:
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    if final_health_status != 503:
+        pytest.skip("issue#3050")
+    assert final_health_status == 503, (
+        "[process_kill health] expected /health 503 after fatal fault, "
+        f"got status={final_health_status}, last_observation={last_observation}"
+    )
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(reason="issue#2327")
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_video_oom_recovers_after_fault_removed(
+    omni_server_function,
+) -> None:
+    """Black-box: after removing OOM pressure, service converges to a terminal state and requests do not hang."""
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    create_url = f"http://{host}:{port}/v1/videos"
+
+    failure_observed = False
+    try:
+        for _ in range(3):
+            try:
+                response = requests.post(
+                    create_url,
+                    data={
+                        "prompt": "oom recover probe",
+                        "width": "512",
+                        "height": "512",
+                        "fps": "8",
+                        "num_frames": "8",
+                        "num_inference_steps": "4",
+                    },
+                    headers={"Accept": "application/json"},
+                    timeout=25,
+                )
+                if response.status_code >= 500:
+                    failure_observed = True
+                    break
+            except Exception:
+                failure_observed = True
+                break
+            time.sleep(1.0)
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+    assert failure_observed, "expected at least one video request failure while OOM pressure is active"
+
+    recovery_deadline = time.monotonic() + 90.0
+    terminal_health: int | None = None
+    unreachable_streak = 0
+    last_health_exc: BaseException | None = None
+    while time.monotonic() < recovery_deadline:
+        try:
+            status, _ = get_health_raw(host, port, timeout_sec=5)
+            unreachable_streak = 0
+            if status in (200, 503):
+                terminal_health = status
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_health_exc = exc
+            msg = str(exc).lower()
+            if "connection refused" in msg or "actively refused" in msg:
+                unreachable_streak += 1
+                if unreachable_streak >= 5:
+                    pytest.fail(
+                        "after OOM sidecar stopped, /health is unreachable repeatedly; "
+                        "the APIServer process likely exited unexpectedly "
+                        f"(last_exc={last_health_exc!r})"
+                    )
+            else:
+                unreachable_streak = 0
+        time.sleep(1.0)
+    else:
+        pytest.fail(
+            "wan22 server did not converge to terminal health (200/503) after OOM pressure was removed; "
+            f"last_health_exc={last_health_exc!r}"
+        )
+
+    assert terminal_health is not None
+    start = time.monotonic()
+    recovery_resp: requests.Response | None = None
+    try:
+        recovery_resp = requests.post(
+            create_url,
+            data={
+                "prompt": "post-recovery admission check",
+                "width": "512",
+                "height": "512",
+                "fps": "8",
+                "num_frames": "8",
+                "num_inference_steps": "4",
+            },
+            headers={"Accept": "application/json"},
+            timeout=30,
+        )
+        recovery_status = recovery_resp.status_code
+        recovery_text_prefix = recovery_resp.text[:300]
+    except Exception as exc:  # noqa: BLE001
+        recovery_status = None
+        recovery_text_prefix = repr(exc)
+    elapsed = time.monotonic() - start
+    assert elapsed < 30, f"post-fault /v1/videos should not hang after OOM removal: {elapsed:.2f}s"
+
+    if terminal_health == 200:
+        assert recovery_status == 200, (
+            "health recovered but /v1/videos admission did not succeed; "
+            f"status={recovery_status}, body={recovery_text_prefix!r}"
+        )
+        assert recovery_resp is not None
+        recovery_payload = recovery_resp.json()
+        assert "id" in recovery_payload, f"post-recovery create payload missing id: {recovery_payload!r}"
+    else:
+        assert recovery_status is None or recovery_status >= 500, (
+            "unhealthy terminal state should fail fast on /v1/videos, "
+            f"got health={terminal_health}, request_status={recovery_status}, body={recovery_text_prefix!r}"
+        )

tests/diffusion/cache/test_teacache_extractors.py

@@ -13,6 +13,8 @@ used by TeaCache. Each model's extractor can be tested by:
 
 Currently implemented:
 - TestFlux2KleinExtractor: Flux2Klein model extractor
+- TestFlux2Extractor: Flux2 model extractor
+- TestFluxExtractor: Flux model extractor
 """
 
 from abc import ABC, abstractmethod
@@ -22,7 +24,12 @@ import pytest
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context
+from vllm_omni.diffusion.cache.teacache.extractors import (
+    extract_flux2_context,
+    extract_flux2_klein_context,
+    extract_flux_context,
+)
+from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -234,7 +241,7 @@ class TestFlux2Extractor(BaseExtractorTest):
         in_channels (128) to inner_dim (6144), so modulated_input should match
         the projected shape, not the input shape.
         """
-        context = extract_flux2_klein_context(flux2_module, **sample_inputs)
+        context = extract_flux2_context(flux2_module, **sample_inputs)
 
         batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
         inner_dim = flux2_module.inner_dim
@@ -277,3 +284,139 @@ class TestFlux2Extractor(BaseExtractorTest):
                 img_ids=torch.randint(0, 64, (1, 1024, 4)),
                 txt_ids=torch.randint(0, 64, (1, 512, 4)),
             )
+
+
+@pytest.mark.cpu
+class TestFluxExtractor(BaseExtractorTest):
+    """Test extract_flux_context function."""
+
+    @pytest.fixture(autouse=True)
+    def cpu_vllm_config(self):
+        """Force CPU custom-op dispatch for this test class."""
+        from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+        with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+            yield
+
+    @pytest.fixture(autouse=True)
+    def mock_flux_attention_backend(self):
+        """Use the SDPA backend so FLUX can be instantiated in CPU tests."""
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+
+        with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend):
+            yield
+
+    def get_extractor(self):
+        return extract_flux_context
+
+    @pytest.fixture
+    def flux_module(self):
+        """Create a minimal FluxTransformer2DModel for testing."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    @pytest.fixture
+    def flux_module_without_guidance(self):
+        """Create a minimal non-guidance-distilled FLUX transformer."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            guidance_embeds=False,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    def get_module(self, flux_module):
+        return flux_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        """Create sample input tensors for Flux."""
+        batch_size = 1
+        img_seq_len = 16
+        txt_seq_len = 8
+        in_channels = 64  # Flux default in_channels
+        txt_dim = 32
+        pooled_dim = 16
+
+        return {
+            "hidden_states": torch.randn(batch_size, img_seq_len, in_channels),
+            "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim),
+            "pooled_projections": torch.randn(batch_size, pooled_dim),
+            "timestep": torch.tensor([500]),
+            "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 3)),
+            "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 3)),
+            "guidance": torch.tensor([3.5]),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_modulated_input_shape(self, flux_module, sample_inputs):
+        """Test that modulated_input has the projected FLUX inner dimension."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+
+        batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
+        assert context.modulated_input.shape == (batch_size, img_seq_len, flux_module.inner_dim)
+
+    def test_run_transformer_blocks_callable(self, flux_module, sample_inputs):
+        """Test that run_transformer_blocks is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.run_transformer_blocks)
+
+    def test_postprocess_callable(self, flux_module, sample_inputs):
+        """Test that postprocess is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.postprocess)
+
+    def test_postprocess_output_shape(self, flux_module, sample_inputs):
+        """Test that postprocess projects back to the input channel width."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        output = context.postprocess(context.hidden_states)
+
+        assert output.sample.shape == sample_inputs["hidden_states"].shape
+
+    def test_postprocess_return_tuple_when_return_dict_false(self, flux_module, sample_inputs):
+        """Test that postprocess honors return_dict=False."""
+        context = extract_flux_context(flux_module, **sample_inputs, return_dict=False)
+        output = context.postprocess(context.hidden_states)
+
+        assert isinstance(output, tuple)
+        assert len(output) == 1
+        assert output[0].shape == sample_inputs["hidden_states"].shape
+
+    def test_without_guidance(self, flux_module_without_guidance, sample_inputs):
+        """Test context extraction works for FLUX variants without guidance embeddings."""
+        inputs = sample_inputs.copy()
+        inputs["guidance"] = None
+
+        context = extract_flux_context(flux_module_without_guidance, **inputs)
+
+        assert context is not None
+        assert context.temb is not None
+
+    def test_invalid_module_raises_error(self):
+        """Test that invalid module without transformer_blocks raises ValueError."""
+        invalid_module = Mock()
+        invalid_module.transformer_blocks = []
+
+        with pytest.raises(ValueError, match="Module must have transformer_blocks"):
+            extract_flux_context(
+                invalid_module,
+                hidden_states=torch.randn(1, 16, 64),
+                encoder_hidden_states=torch.randn(1, 8, 32),
+                pooled_projections=torch.randn(1, 16),
+                timestep=torch.tensor([500]),
+                img_ids=torch.randint(0, 64, (1, 16, 3)),
+                txt_ids=torch.randint(0, 64, (1, 8, 3)),
+            )

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -121,7 +121,8 @@ class TestTrajectoryRecording:
 
         assert trajectory_latents is not None
         assert trajectory_timesteps is not None
-        assert len(trajectory_latents) == EXPECTED_STEPS
+        # initial latent + one per denoising step
+        assert len(trajectory_latents) == EXPECTED_STEPS + 1
         assert len(trajectory_timesteps) == EXPECTED_STEPS
         # log_probs is None without a scheduler (default ODE path)
         assert trajectory_log_probs is None
@@ -167,6 +168,54 @@ class TestTrajectoryRecording:
             "Last trajectory latent should match the final output"
         )
 
+    def test_initial_latent_matches_input_noise(self, bagel_and_args):
+        """Regression: the first trajectory entry must be the initial noise (pre-denoising)."""
+        bagel, args = bagel_and_args
+        init_noise = args["packed_init_noises"].clone()
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert torch.allclose(trajectory_latents[0], init_noise, atol=1e-6), (
+            "trajectory_latents[0] should be the initial noise before any denoising step"
+        )
+
+    def test_trajectory_timesteps_match_expected_schedule(self, bagel_and_args):
+        """Regression: trajectory_timesteps must record raw timesteps[i], not timesteps[i] - dts[i]."""
+        bagel, args = bagel_and_args
+
+        # Recompute the expected timestep schedule (mirrors generate_image logic)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
+        shift = args.get("timestep_shift", 1.0)
+        ts = shift * ts / (1 + (shift - 1) * ts)
+        expected_timesteps = ts[:-1]  # last element is dropped
+
+        _, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert len(trajectory_timesteps) == len(expected_timesteps)
+        for i, (actual, expected) in enumerate(zip(trajectory_timesteps, expected_timesteps)):
+            assert abs(float(actual) - float(expected)) < 1e-6, (
+                f"Step {i}: trajectory_timestep={float(actual):.6f}, expected={float(expected):.6f}"
+            )
+
+    def test_latent_count_equals_timesteps_plus_one(self, bagel_and_args):
+        """Regression: len(trajectory_latents) == len(timesteps) + 1 (initial + one per step)."""
+        bagel, args = bagel_and_args
+
+        # Compute the number of denoising steps the same way generate_image does
+        ts = torch.linspace(1, 0, args["num_timesteps"])
+        shift = args.get("timestep_shift", 1.0)
+        ts = shift * ts / (1 + (shift - 1) * ts)
+        num_steps = len(ts) - 1  # timesteps = ts[:-1]
+
+        _, trajectory_latents, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        assert len(trajectory_latents) == num_steps + 1, (
+            f"Expected {num_steps + 1} latents (initial + {num_steps} steps), got {len(trajectory_latents)}"
+        )
+        assert len(trajectory_timesteps) == num_steps, (
+            f"Expected {num_steps} timesteps, got {len(trajectory_timesteps)}"
+        )
+
 
 # ---------------------------------------------------------------------------
 # Mock scheduler for log-prob tests

tests/diffusion/test_diffusion_engine_metrics.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Source-level regression tests for DiffusionEngine.
+
+These tests verify naming conventions and patterns by inspecting source code
+at the function level using AST. They are intentionally coupled to the source
+layout and should be updated whenever the metrics construction code is
+refactored.
+"""
+
+from __future__ import annotations
+
+import ast
+import os
+
+_ENGINE_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "diffusion",
+        "diffusion_engine.py",
+    )
+)
+
+
+def _read_engine_source() -> str:
+    with open(_ENGINE_PATH) as f:
+        return f.read()
+
+
+def _get_function_source(source: str, class_name: str | None, func_name: str) -> str:
+    """Extract the source of a specific function/method using AST.
+
+    Args:
+        source: Full file source code.
+        class_name: Enclosing class name, or None for module-level functions.
+        func_name: Function/method name.
+
+    Returns:
+        Source code of the function body.
+    """
+    tree = ast.parse(source)
+    if class_name is not None:
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef) and node.name == class_name:
+                for item in node.body:
+                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name:
+                        result = ast.get_source_segment(source, item)
+                        assert result is not None, f"{class_name}.{func_name} source not found"
+                        return result
+    else:
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
+                result = ast.get_source_segment(source, node)
+                assert result is not None, f"{func_name} source not found"
+                return result
+    raise AssertionError(f"Function {class_name + '.' if class_name else ''}{func_name} not found in source")
+
+
+class TestMetricKeys:
+    """Verify metric naming conventions in DiffusionEngine.step()."""
+
+    def test_no_duplicate_preprocess_key(self) -> None:
+        """step() should not contain 'preprocessing_time_ms' (duplicate of 'preprocess_time_ms')."""
+        source = _read_engine_source()
+        step_source = _get_function_source(source, "DiffusionEngine", "step")
+        assert "preprocessing_time_ms" not in step_source, (
+            "Found duplicate key 'preprocessing_time_ms' in step() — should only use 'preprocess_time_ms'"
+        )
+
+    def test_metric_key_naming_consistency(self) -> None:
+        """In step(): exec_time should use exec_total_time, total_time should use step_total_ms."""
+        source = _read_engine_source()
+        step_source = _get_function_source(source, "DiffusionEngine", "step")
+        lines = step_source.split("\n")
+
+        found_exec = False
+        found_total = False
+        for line in lines:
+            if '"diffusion_engine_exec_time_ms"' in line:
+                found_exec = True
+                assert "exec_total_time" in line, (
+                    "diffusion_engine_exec_time_ms should measure executor time only (exec_total_time)"
+                )
+            if '"diffusion_engine_total_time_ms"' in line:
+                found_total = True
+                assert "step_total_ms" in line, (
+                    "diffusion_engine_total_time_ms should measure full step time (step_total_ms)"
+                )
+        assert found_exec, "diffusion_engine_exec_time_ms key not found in step()"
+        assert found_total, "diffusion_engine_total_time_ms key not found in step()"
+
+
+class TestDummyRunAllocation:
+    """Verify _dummy_run generates exact-sized audio arrays."""
+
+    def test_no_oversized_allocation(self) -> None:
+        """_dummy_run should not allocate more audio than needed."""
+        source = _read_engine_source()
+        dummy_source = _get_function_source(source, "DiffusionEngine", "_dummy_run")
+        assert "audio_sr * audio_duration_sec" not in dummy_source, (
+            "_dummy_run should generate exact-sized audio, not allocate and slice"
+        )

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -12,6 +12,7 @@ from pytest_mock import MockerFixture
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
@@ -111,7 +112,11 @@ def test_load_poll(build_adapter):
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.load_async(request)
-    payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True}
+    payload: OmniPayload = {
+        "codes": {"audio": [[1]]},
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 16)
     adapter._poll_single_request(request)
 
@@ -154,12 +159,49 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
-    adapter._update_request_payload("ext", {"h": torch.tensor([[1.0]]), "codes": [1], "finished": False})
-    merged = adapter._update_request_payload("ext", {"h": torch.tensor([[2.0]]), "codes": [2], "finished": True})
+    first: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "codes": {"audio": [1]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    adapter._update_request_payload("ext", first)
+    second: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "codes": {"audio": [2]},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    merged = adapter._update_request_payload("ext", second)
 
-    assert torch.equal(merged["h"], torch.tensor([[1.0], [2.0]]))
-    assert merged["codes"] == [1, 2]
-    assert merged["finished"] is True
+    assert torch.equal(merged["hidden_states"]["output"], torch.tensor([[1.0], [2.0]]))
+    assert merged["codes"]["audio"] == [1, 2]
+    assert merged["meta"]["finished"].item() is True
+
+
+def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
+
+    adapter.request_ids_mapping["req-merged"] = "ext-merged"
+    adapter.request_payload["ext-merged"] = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "ids": {"prompt": [11, 12]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    connector.get.return_value = (payload, 8)
+
+    adapter._poll_single_request(request)
+
+    assert torch.equal(
+        request.additional_information["hidden_states"]["output"],
+        torch.tensor([[1.0], [2.0]]),
+    )
+    # Keys absent from the new chunk are dropped (matches main's behavior).
+    assert "ids" not in request.additional_information
+    assert request.additional_information["meta"]["finished"].item() is True
 
 
 def test_process_and_restore_queues(build_adapter):
@@ -321,7 +363,10 @@ def test_cleanup_after_poll_flow(build_adapter):
     adapter.load_async(request)
 
     adapter.request_ids_mapping["req-flow"] = "ext-flow"
-    payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True}
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 8)
     adapter._poll_single_request(request)
 
@@ -411,8 +456,8 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         client_index=0,
         take_events=lambda: [],
         trace_headers=None,
-        num_cached_tokens=0,
-        num_external_computed_tokens=0,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
         num_nans_in_logits=0,
         get_finished_reason=lambda: "stop",
     )
@@ -492,8 +537,8 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         client_index=0,
         take_events=lambda: [],
         trace_headers=None,
-        num_cached_tokens=0,
-        num_external_computed_tokens=0,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
         num_nans_in_logits=0,
         get_finished_reason=lambda: "stop",
     )

tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py

@@ -17,6 +17,7 @@ import time
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
     ManagedBuffer,
     MooncakeTransferEngineConnector,
@@ -24,7 +25,7 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
 )
 
 # All tests in this file require Mooncake TransferEngine and an RDMA environment.
-pytestmark = [pytest.mark.parallel, pytest.mark.gpu]
+pytestmark = [pytest.mark.parallel, pytest.mark.core_model]
 
 # ---------------------------------------------------------------------------
 # Shared helpers
@@ -194,7 +195,7 @@ def _md5(tensor: torch.Tensor) -> str:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestBasicConnector:
     """Verify connector initialization, put, cleanup, and health check."""
 
@@ -254,7 +255,7 @@ class TestBasicConnector:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestEndToEnd:
     """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
 
@@ -447,7 +448,7 @@ class TestEndToEnd:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestLifecycle:
     """Close, context manager, double-close safety."""
 
@@ -475,9 +476,7 @@ class TestLifecycle:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.cuda
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestGPUPool:
     """GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
 
@@ -521,12 +520,11 @@ class TestGPUPool:
 
 
 # ---------------------------------------------------------------------------
-# 5. Stress / Correctness tests (marked slow, skipped in quick CI)
+# 5. Stress / Correctness tests, slow
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.slow
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestStressCorrectness:
     """
     Slow but high-value regression tests: concurrent put+get with data

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -67,7 +67,7 @@ if current_omni_platform.is_rocm():
     ]
 
 # Maximum allowed difference per color channel
-PIXEL_TOLERANCE = 5
+PIXEL_TOLERANCE = 10
 
 # Default test prompt
 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
@@ -304,11 +304,13 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
     if not _is_mooncake_master_available():
-        pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
+        raise RuntimeError(
+            "mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)"
+        )
     MOONCAKE_HOST = "127.0.0.1"
     MOONCAKE_RPC_PORT = _find_free_port()
     MOONCAKE_HTTP_PORT = _find_free_port()

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -13,10 +13,8 @@ verifies a stable reference recipe:
 from __future__ import annotations
 
 import functools
-import io
 import os
 from pathlib import Path
-from urllib.request import urlopen
 
 import numpy as np
 import pytest
@@ -25,6 +23,7 @@ from huggingface_hub import snapshot_download
 from vllm.sampling_params import SamplingParams
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import test_asset_path
 from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
@@ -33,7 +32,9 @@ from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokeni
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"
 
-REFERENCE_PROMPT_WAV_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+# Vendored under tests/assets/cosyvoice3/ so the test does not depend on
+# raw.githubusercontent.com being reachable from CI runners.
+REFERENCE_PROMPT_WAV_PATH = test_asset_path("cosyvoice3/zero_shot_prompt.wav")
 REFERENCE_PROMPT_TEXT = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
 REFERENCE_SYNTH_TEXT = (
     "CosyVoice is undergoing a comprehensive upgrade, providing more accurate, "
@@ -51,9 +52,7 @@ ASYNC_CHUNK_MODES = [
 
 @functools.lru_cache(maxsize=1)
 def _load_reference_prompt_wav() -> tuple[np.ndarray, int]:
-    with urlopen(REFERENCE_PROMPT_WAV_URL, timeout=30) as resp:
-        data = resp.read()
-    audio, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
+    audio, sr = sf.read(str(REFERENCE_PROMPT_WAV_PATH), dtype="float32", always_2d=False)
     if isinstance(audio, np.ndarray) and audio.ndim > 1:
         audio = np.mean(audio, axis=-1)
     return np.asarray(audio, dtype=np.float32), int(sr)