fix (#2946)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

.buildkite/pipeline-npu.yaml

@@ -0,0 +1,92 @@
+steps:
+  - label: ":buildkit: Build and Push NPU Test Image"
+    key: image-build
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-2"
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              vault.hashicorp.com/agent-init-first: "true"
+              vault.hashicorp.com/agent-inject: "true"
+              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
+              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
+              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
+              vault.hashicorp.com/agent-pre-populate-only: "true"
+              vault.hashicorp.com/agent-run-as-group: "1000"
+              vault.hashicorp.com/agent-run-as-user: "1000"
+              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
+              vault.hashicorp.com/role: ascend-gha-runners
+              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
+              vault.hashicorp.com/tls-skip-verify: "true"
+          podSpecPatch:
+            volumes:
+            - name: token-vol
+              projected:
+                defaultMode: 420
+                sources:
+                - serviceAccountToken:
+                    audience: api
+                    expirationSeconds: 600
+                    path: token
+    env:
+      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+      IMAGE_NAME: "vllm-omni-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: |
+      set -ex
+
+      echo "--- Building and pushing NPU Test Image"
+      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+      echo "buildkitd address: $${BUILDKITD_ADDR}"
+
+      git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
+
+      if ! command -v buildctl &> /dev/null; then
+        echo "Installing buildctl..."
+        mkdir -p /tmp/buildkit
+        BUILDKIT_VERSION="v0.29.0"
+        wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
+        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+      fi
+
+      export DOCKER_CONFIG=/home/user/.docker
+      buildctl \
+        --addr="$${BUILDKITD_ADDR}" \
+        --tlscacert=/home/user/.docker/ca.pem \
+        --tlscert=/home/user/.docker/cert.pem \
+        --tlskey=/home/user/.docker/key.pem \
+        build \
+        --frontend dockerfile.v0 \
+        --local context=. \
+        --local dockerfile=./docker \
+        --opt filename=Dockerfile.npu \
+        --secret id=dockerconfig,src=/home/user/.docker/config.json \
+        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
+        --progress=plain
+
+      echo "--- Image pushed successfully"
+      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    if: build.branch != "main" && build.pull_request.labels includes "ready"
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-2"
+    command: |
+      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
+      

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

@@ -64,15 +64,22 @@ while true; do
 done
 
 echo "--- Pulling container"
-image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+## Temporary change to use AMD Docker Hub to store the vllm-ci image
+# to bypass the rate limit issue with ECR Public Gallery.
+# TODO: @tjtanaa point back to ECR Public Gallery
+# once the amd agents are configured to use ECR Public Gallery.
+# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 
-# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
-sudo apt-get update && sudo apt-get install -y awscli
-# Use safe docker login helper to prevent race conditions
-source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
-safe_docker_login_ecr_public
-# Pull the container from ECR Public Gallery
+# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
+# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+# sudo apt-get update && sudo apt-get install -y awscli
+##  Use safe docker login helper to prevent race conditions
+# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+# safe_docker_login_ecr_public
+
+## Pull the container from AMD Docker Hub
 
 docker pull "${image_name}"
 
@@ -126,8 +133,6 @@ if [[ $commands == *"--shard-id="* ]]; then
         -e VLLM_ROCM_USE_AITER=0 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
-        -e AWS_ACCESS_KEY_ID \
-        -e AWS_SECRET_ACCESS_KEY \
         -v "${HF_CACHE}:${HF_MOUNT}" \
         -e "HF_HOME=${HF_MOUNT}" \
         -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -160,8 +165,6 @@ else
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
           -e VLLM_ROCM_USE_AITER=0 \
           -e HF_TOKEN \
-          -e AWS_ACCESS_KEY_ID \
-          -e AWS_SECRET_ACCESS_KEY \
           -v "${HF_CACHE}:${HF_MOUNT}" \
           -e "HF_HOME=${HF_MOUNT}" \
           -e "PYTHONPATH=${MYPYTHONPATH}" \

.buildkite/test-npu-ready.yaml

@@ -0,0 +1,14 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-ready-pipeline
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-2"
+    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
+    plugins:
+      - kubernetes:
+          podSpecPatch:
+            imagePullSecrets:
+            - name: swr-secret
+    commands:
+      - "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"

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

@@ -3,7 +3,7 @@
    Last synced: 2025-12-15
    Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
   - group: "AMD Tests"
@@ -13,7 +13,6 @@
         depends_on: ~
         soft_fail: false
         commands:
-          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
           - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
@@ -22,13 +21,15 @@
         retry:
           automatic:
             - exit_status: -1  # Agent was lost
-              limit: 1
+              limit: 2
             - exit_status: -10  # Agent was lost
-              limit: 1
+              limit: 2
+            - exit_status: 128  # Git connectivity issues
+              limit: 2
             - exit_status: 1  # Machine occasionally fail
               limit: 1
         agents:
-          queue: cpu_queue_premerge
+          queue: amd-cpu
 
     {% for step in steps %}
     {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}

docker/Dockerfile.npu

@@ -5,26 +5,34 @@ FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 WORKDIR /vllm-workspace/vllm
 RUN git fetch origin --tags && git checkout v0.18.0
 
+RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
+
 WORKDIR /vllm-workspace/vllm-ascend
 RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 # Install vllm-ascend
-# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --extra-index https://download.pytorch.org/whl/cpu/ \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -5,26 +5,34 @@ FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 WORKDIR /vllm-workspace/vllm
 RUN git fetch origin --tags && git checkout v0.18.0
 
+RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
+
 WORKDIR /vllm-workspace/vllm-ascend
 RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 # Install vllm-ascend
-# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --extra-index https://download.pytorch.org/whl/cpu/ \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+    python3 -m pip install -U pip && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
+      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
+      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
+      --default-timeout=60
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docs/contributing/ci/CI_5levels.md

@@ -567,15 +567,15 @@ L4 level testing is a comprehensive quality audit before a version release. It e
             {
                 "dataset_name": "random",
                 "num_prompts": [10, 20],
-                "request_rate": [0.5, 1],
+                "max_concurrency": [1, 4],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [500, 800],
+                    "mean_audio_ttfp_ms": [2000, 3500],
+                    "mean_audio_rtf": [0.25, 0.35]
                 }
             }
         ]
@@ -634,15 +634,15 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 
     1.  Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
     2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
-    3.  Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
+    3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
     4.  The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
 
     | Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
     | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
     | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
     | request_rate    | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-    | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-    </details>
+    | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
+    | baseline        | object      | No       | see above       | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).                                                                                                               |
 
     -   -   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
 

docs/contributing/ci/failures.md

@@ -1,10 +1,106 @@
 # CI Failures
 
-What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
+## Overview of CI Checks
 
+When you open a PR against vLLM-Omni, several CI checks run automatically:
 
-## Common Case 1:
+| Check | Platform | What it does |
+| ----- | -------- | ------------ |
+| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
+| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
+| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
+| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
+| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
+| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
+| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |
 
-🚨 Error: The command was interrupted by a signal: signal: terminated
+## Step 1: Identify the Failing Check
 
-Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
+Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
+
+### pre-commit failures
+
+These are typically formatting or linting issues introduced by your PR. Fix them locally:
+
+```bash
+uv pip install pre-commit
+pre-commit run --all-files
+```
+
+Then commit the fixes and push.
+
+### DCO failures
+
+Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
+
+```bash
+git commit --amend -s
+git push --force-with-lease
+```
+
+For multiple commits, use an interactive rebase to add the sign-off to each one.
+
+### Read the Docs failures
+
+The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
+
+```bash
+uv pip install -e ".[docs]"
+mkdocs build --strict
+```
+
+Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
+
+### Buildkite failures
+
+Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
+
+## Step 2: Check if the Failure Is a Known Issue
+
+Before spending time debugging, check whether the failure already exists on the `main` branch:
+
+1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
+2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
+3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
+
+If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
+
+## Step 3: Investigate the Failure
+
+If the failure appears to be new, investigate whether your changes caused it.
+
+### Reading Buildkite Logs
+
+1. Click **Details** next to the Buildkite check on your PR.
+2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
+3. Expand the step to see the full test output with the traceback.
+
+### Running Tests Locally
+
+For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
+
+## Step 4: Raise an Issue or Fix It
+
+### If the failure is pre-existing (not caused by your PR)
+
+1. **Raise a new issue** if one doesn't already exist, using the title format:
+   `[CI Failure]: [job-name] - [test-path]`
+2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
+3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
+
+### If the failure is caused by your PR
+
+1. Fix the issue in your branch and push the update.
+2. If the fix is non-trivial, consider adding a test to prevent regression.
+
+## Common Failure Patterns
+
+| Symptom | Likely Cause | Fix |
+| ------- | ------------ | --- |
+| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
+| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
+| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
+| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
+| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
+| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
+| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |

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

@@ -29,6 +29,11 @@ uv pip install vllm --torch-backend=auto
 uv pip install vllm-omni
 ```
 
+To run Gradio demos, also install the optional extras:
+```bash
+uv pip install 'vllm-omni[demo]'
+```
+
 # --8<-- [end:pre-built-wheels]
 
 # --8<-- [start:build-wheel-from-source]
@@ -50,6 +55,11 @@ cd vllm-omni
 uv pip install -e .
 ```
 
+To run Gradio demos, install with optional extras:
+```bash
+uv pip install -e '.[demo]'
+```
+
 <details><summary>(Optional) Installation of vLLM from source</summary>
 If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
 

docs/user_guide/diffusion/quantization/fp8.md

@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | 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 |
+| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion/teacache.md

@@ -125,6 +125,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
+| `HunyuanImage3Pipeline` | HunyuanImage3 | `tencent/HunyuanImage-3.0-Instruct` |
 
 ### VideoGen
 

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -256,6 +256,15 @@ lora_adapter/
 
 ## Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Launch the Gradio demo:
 
 ```bash

docs/user_guide/examples/online_serving/fish_speech.md

@@ -11,6 +11,15 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serv
 
 ## Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
 
 ```bash

docs/user_guide/examples/online_serving/image_to_image.md

@@ -122,6 +122,15 @@ python openai_chat_client.py --input input1.png input2.png --prompt "Combine the
 
 ### Method 4: Using Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7861

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -149,6 +149,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 
 ## Run Local Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -154,7 +154,7 @@ print(response.choices[1].message.audio)    # Audio response
 ```
 
 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py \
     --query-type use_image \
@@ -163,6 +163,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 
 ## Run Local Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -19,6 +19,15 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 
 ## Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Two interactive Gradio demos are available, both supporting all 3 task types:
 
 | Demo     | File                     | Transport    | Streaming Quality                                  |

docs/user_guide/examples/online_serving/text_to_image.md

@@ -105,6 +105,15 @@ python openai_chat_client.py --prompt "A beautiful landscape painting" --output
 
 ### Method 4: Using Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7860

examples/offline_inference/glm_image/README.md

@@ -129,8 +129,17 @@ The first run loads model weights. Subsequent runs are faster:
 --stage-init-timeout 900  # Increase timeout for slow storage
 ```
 
+### `Transformers does not recognize this architecture` Error
+
+Your have to upgrade `transformers` package to `5.3.0` or above:
+
+```
+pip install --upgrade transformers
+```
+
 ## Requirements
 
 - vLLM-Omni with GLM-Image support
 - CUDA-capable GPU (recommended: H100/A100 with 80GB)
 - GLM-Image model weights
+- `transformers` v5.3.0 or above

examples/offline_inference/helios/end2end.py

@@ -178,6 +178,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8').",
+    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -213,6 +219,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        quantization=args.quantization,
     )
 
     # Validate I2V / V2V arguments

examples/offline_inference/text_to_image/gradio_demo.py

@@ -1,7 +1,10 @@
 import argparse
 from functools import lru_cache
 
-import gradio as gr
+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 torch
 
 from vllm_omni.entrypoints.omni import Omni

examples/offline_inference/text_to_image/text_to_image.py

@@ -64,8 +64,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--guidance-scale",
         type=float,
-        default=1.0,
-        help="Classifier-free guidance scale.",
+        default=4.0,
+        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
     )
     parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
     parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")

examples/online_serving/fish_speech/gradio_demo.py

@@ -16,7 +16,10 @@ import argparse
 import base64
 import io
 
-import gradio as gr
+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

examples/online_serving/glm_image/README.md

@@ -196,3 +196,9 @@ gpu_memory_utilization: 0.5
 ```
 
 - The first request may be slow due to model warmup. Subsequent requests will be faster.
+
+- If you encounter `Transformers does not recognize this architecture` error, your have to upgrade `transformers` package to `5.3.0` or above:
+
+```
+pip install --upgrade transformers
+```

examples/online_serving/image_to_image/gradio_demo.py

@@ -10,7 +10,10 @@ import argparse
 import base64
 from io import BytesIO
 
-import gradio as gr
+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 requests
 from PIL import Image
 

examples/online_serving/qwen2_5_omni/gradio_demo.py

@@ -6,7 +6,10 @@ import random
 from pathlib import Path
 from typing import Any
 
-import gradio as gr
+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 numpy as np
 import soundfile as sf
 import torch

examples/online_serving/qwen3_omni/README.md

@@ -211,7 +211,7 @@ print(response.choices[1].message.audio)
 Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default.
 
 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
 ```bash
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
     --query-type use_image \

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -6,7 +6,10 @@ import random
 from pathlib import Path
 from typing import Any
 
-import gradio as gr
+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 numpy as np
 import soundfile as sf
 import torch

examples/online_serving/qwen3_tts/gradio_demo.py

@@ -24,7 +24,10 @@ import io
 import json
 import logging
 
-import gradio as gr
+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

examples/online_serving/qwen3_tts/tts_common.py

@@ -3,7 +3,10 @@
 import base64
 import io
 
-import gradio as gr
+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

examples/online_serving/text_to_image/gradio_demo.py

@@ -10,7 +10,10 @@ import argparse
 import base64
 from io import BytesIO
 
-import gradio as gr
+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 requests
 from PIL import Image
 

examples/online_serving/voxtral_tts/gradio_demo.py

@@ -20,7 +20,10 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any
 
-import gradio as gr
+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

pyproject.toml

@@ -56,6 +56,10 @@ dev = [
     "mistune>=3.2.0", # for example tests
 ]
 
+demo = [
+    "gradio>=6.7.0",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

requirements/common.txt

@@ -4,7 +4,6 @@ librosa>=0.11.0
 resampy>=0.4.3
 diffusers>=0.36.0
 accelerate==1.12.0
-gradio==5.50
 soundfile>=0.13.1
 cache-dit==1.3.0
 tqdm>=4.66.0
@@ -18,3 +17,4 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
+pydub

tests/dfx/perf/scripts/run_benchmark.py

@@ -118,10 +118,6 @@ def benchmark_params(request, omni_server):
     if param_index >= len(all_params):
         raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")
 
-    if all_params[param_index]["dataset_name"] == "random-mm":
-        # TODO: Due to known issues, skip the random-mm dataset.
-        pytest.skip("Skipping parameter for random-mm dataset.")
-
     current = param_index + 1
     total = len(all_params)
     print(f"\n  Running benchmark {current}/{total} for {test_name}")
@@ -132,15 +128,70 @@ def benchmark_params(request, omni_server):
     }
 
 
-def assert_result(result, params, num_prompt):
+def _resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> Any:
+    """Pick the baseline threshold for this sweep step.
+
+    Supported shapes per metric:
+    - **Scalar** — same threshold for every concurrency / QPS.
+    - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
+    - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
+
+    For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
+    """
+    if baseline_raw is None:
+        # If no baseline is set, the maximum value will be used.
+        return 100000
+    if isinstance(baseline_raw, dict):
+        if max_concurrency is not None:
+            for key in (max_concurrency, str(max_concurrency)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        if request_rate is not None:
+            for key in (request_rate, str(request_rate)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        raise KeyError(
+            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
+            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def assert_result(
+    result,
+    params,
+    num_prompt,
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> None:
     assert result["completed"] == num_prompt, "Request failures exist"
     baseline_data = params.get("baseline", {})
-    for metric_name, baseline_value in baseline_data.items():
+    for metric_name, baseline_raw in baseline_data.items():
         current_value = result[metric_name]
+        baseline_value = _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
         if "throughput" in metric_name:
-            assert current_value >= baseline_value, f"{metric_name}: {current_value} < {baseline_value}"
+            if current_value <= baseline_value:
+                print(
+                    f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
+                )
         else:
-            assert current_value <= baseline_value, f"{metric_name}: {current_value} > {baseline_value}"
+            if current_value >= baseline_value:
+                print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
 
 
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
@@ -195,8 +246,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
         elif not isinstance(value, bool):
             args.extend([arg_name, str(value)])
 
-    # QPS test
-    for qps, num_prompt in zip(qps_list, num_prompt_list):
+    # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
+    for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
             args=args,
@@ -205,10 +256,16 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
         )
-        assert_result(result, params, num_prompt=num_prompt)
+        assert_result(
+            result,
+            params,
+            num_prompt=num_prompt,
+            sweep_index=i,
+            request_rate=qps,
+        )
 
-    # concurrency test
-    for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
+    # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
+    for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
             args=args,
@@ -217,4 +274,10 @@ def test_performance_benchmark(omni_server, benchmark_params):
             dataset_name=dataset_name,
             num_prompt=num_prompt,
         )
-        assert_result(result, params, num_prompt=num_prompt)
+        assert_result(
+            result,
+            params,
+            num_prompt=num_prompt,
+            sweep_index=i,
+            max_concurrency=concurrency,
+        )

tests/dfx/perf/tests/test.json

@@ -25,9 +25,9 @@
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [1000, 3000, 5000],
+                    "mean_audio_ttfp_ms": [8000, 10000, 13000],
+                    "mean_audio_rtf": [0.2, 0.25, 0.45]
                 }
             },
             {
@@ -39,10 +39,10 @@
                     40,
                     100
                 ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    10
+                "request_rate": [
+                    0.1,
+                    0.3,
+                    0.5
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -62,9 +62,31 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [2000, 4000, 6000],
+                    "mean_audio_ttfp_ms": [10000, 13000, 15000],
+                    "mean_audio_rtf": [0.25, 0.35, 0.45]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000],
+                    "mean_audio_ttfp_ms": [30000, 60000],
+                    "mean_audio_rtf": [0.35, 0.45]
                 }
             }
         ]
@@ -113,9 +135,9 @@
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [1000, 3000, 5000],
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000],
+                    "mean_audio_rtf": [0.2, 0.35, 0.6]
                 }
             },
             {
@@ -127,10 +149,10 @@
                     40,
                     100
                 ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    10
+                "request_rate": [
+                    0.1,
+                    0.3,
+                    0.5
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -150,9 +172,31 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [2000, 4000, 6000],
+                    "mean_audio_ttfp_ms": [2000, 4000, 6000],
+                    "mean_audio_rtf": [0.25, 0.4, 0.7]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000],
+                    "mean_audio_ttfp_ms": [1000, 3000],
+                    "mean_audio_rtf": [0.35, 0.45]
                 }
             }
         ]
@@ -183,8 +227,8 @@
                 },
                 "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_audio_ttfp_ms": [6000, 6000],
+                    "mean_audio_rtf": [0.3, 0.3]
                 }
             }
         ]

tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py

@@ -7,6 +7,49 @@ import pytest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+class TestSetForwardContextNumTokens:
+    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
+
+    def test_sets_num_tokens_when_context_available(self, mocker):
+        """num_tokens should be set on ForwardContext when available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        del mock_ctx.in_profile_run  # simulate missing attr
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(1024)
+
+        assert mock_ctx.num_tokens == 1024
+        assert mock_ctx.in_profile_run is False
+
+    def test_sets_in_profile_run_only_if_missing(self, mocker):
+        """in_profile_run should not be overwritten if already set."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.in_profile_run = True  # already set
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(512)
+
+        assert mock_ctx.num_tokens == 512
+        assert mock_ctx.in_profile_run is True  # not overwritten
+
+    def test_noop_when_context_unavailable(self, mocker):
+        """Should do nothing when ForwardContext is not available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
+        mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
+
+        hunyuan_moe._set_forward_context_num_tokens(256)
+
+        mock_get.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

tests/entrypoints/openai_api/test_image_server.py

@@ -115,7 +115,7 @@ class MockGenerationResult:
 class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
-    def __init__(self):
+    def __init__(self, images=None):
         self.stage_configs = [
             SimpleNamespace(stage_type="llm"),
             SimpleNamespace(stage_type="diffusion"),
@@ -123,11 +123,12 @@ class FakeAsyncOmni:
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
         self.captured_prompt = None
+        self._images = images or [Image.new("RGB", (64, 64), color="green")]
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_sampling_params_list = sampling_params_list
         self.captured_prompt = prompt
-        images = [Image.new("RGB", (64, 64), color="green")]
+        images = [img.copy() for img in self._images]
         yield MockGenerationResult(images)
 
 
@@ -168,7 +169,13 @@ def test_client(mock_async_diffusion):
     app.state.engine_client = mock_async_diffusion
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
-    app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
+    from vllm.entrypoints.openai.models.protocol import BaseModelPath
+
+    from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels
+
+    app.state.openai_serving_models = _DiffusionServingModels(
+        [BaseModelPath(name="Qwen/Qwen-Image", model_path="Qwen/Qwen-Image")]
+    )
     app.state.args = Namespace(
         default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
         max_generated_image_size=4096,  # 64*64
@@ -199,6 +206,28 @@ def async_omni_test_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def async_omni_rgba_test_client():
+    """Create test client with mocked AsyncOmni engine returning RGBA output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    app.state.engine_client = FakeAsyncOmni(images=[Image.new("RGBA", (64, 64), color=(0, 255, 0, 128))])
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=1048576,
+    )
+    return TestClient(app)
+
+
 @pytest.fixture
 def async_omni_stage_configs_only_client():
     """Create test client with refactored AsyncOmni compatibility surface only."""
@@ -977,6 +1006,27 @@ def test_image_edit_compression_jpeg(test_client):
     assert len(img_bytes_50) < len(img_bytes_100)
 
 
+def test_image_edit_rgba_output_converts_to_jpeg(async_omni_rgba_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    response = async_omni_rgba_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+        },
+    )
+    assert response.status_code == 200
+
+    data = response.json()
+    img_bytes = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes))
+    assert img.format.lower() == "jpeg"
+    assert img.mode == "RGB"
+    assert data["output_format"] == "jpeg"
+
+
 def test_image_edit_compression_png(async_omni_test_client):
     img_bytes_1 = make_test_image_bytes((16, 16))
     # uploadfile with image key

vllm_omni/diffusion/cache/teacache/backend.py

@@ -20,6 +20,24 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 logger = init_logger(__name__)
 
 
+def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """
+    Enable TeaCache for HunyuanImage3 model.
+
+    HunyuanImage3 uses a GPT-based architecture with KV cache, which is incompatible
+    with the standard hook-based TeaCache approach. Instead, we store the TeaCacheConfig
+    on the pipeline so the denoising loop can implement caching directly.
+    """
+    teacache_config = TeaCacheConfig(
+        transformer_type="HunyuanImage3Pipeline",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    pipeline._tea_cache_config = teacache_config
+
+    logger.info(f"TeaCache enabled for HunyuanImage3 with rel_l1_thresh={teacache_config.rel_l1_thresh}")
+
+
 def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """
     Enable TeaCache for Bagel model.
@@ -70,6 +88,7 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) ->
 CUSTOM_TEACACHE_ENABLERS = {
     "BagelPipeline": enable_bagel_teacache,
     "Flux2KleinPipeline": enable_flux2_klein_teacache,
+    "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
 }
 
 
@@ -157,6 +176,17 @@ class TeaCacheBackend(CacheBackend):
                                 Currently not used by TeaCache but accepted for interface consistency.
             verbose: Whether to log refresh operations (default: True)
         """
+        # HunyuanImage3: tea cache state is managed inside the denoising loop,
+        # so refresh is a no-op (state is re-initialized every __call__).
+        if (
+            hasattr(pipeline, "_tea_cache_config")
+            and isinstance(pipeline._tea_cache_config, TeaCacheConfig)
+            and pipeline.__class__.__name__ == "HunyuanImage3Pipeline"
+        ):
+            if verbose:
+                logger.debug(f"TeaCache state refreshed for HunyuanImage3 (num_inference_steps={num_inference_steps})")
+            return
+
         # Extract transformer from pipeline
         transformer = pipeline.transformer
 

vllm_omni/diffusion/cache/teacache/config.py

@@ -54,6 +54,16 @@ _MODEL_COEFFICIENTS = {
         -12.281286412689623,
         1.0733905006198015,
     ],
+    # HunyuanImage3 pipeline coefficients
+    # Calibrated via polyfit on 3920 data points (80 prompts × 49 steps).
+    # Maps time_embed rel_l1 (range ~0.12-0.54) to output rel_l1 (range ~0.01-0.53).
+    "HunyuanImage3Pipeline": [
+        1.04117826e02,
+        -1.26848482e02,
+        5.68168652e01,
+        -1.04182570e01,
+        6.78098549e-01,
+    ],
 }
 
 

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -196,6 +196,8 @@ class FluxAttention(torch.nn.Module):
         attention_mask: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # Ensure contiguous for FP8 quantized linear layers
+        hidden_states = hidden_states.contiguous()
         qkv, _ = self.to_qkv(hidden_states)
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
@@ -209,6 +211,7 @@ class FluxAttention(torch.nn.Module):
         key = self.norm_k(key)
 
         if self.added_kv_proj_dim is not None:
+            encoder_hidden_states = encoder_hidden_states.contiguous()
             encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
             add_q_size = self.add_kv_proj.num_heads * self.head_dim
             add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
@@ -720,6 +723,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
         guidance_embeds: bool = True,
         axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
         theta: float = 10000.0,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__(
             od_config=od_config,
@@ -735,6 +739,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
             guidance_embeds=guidance_embeds,
             axes_dims_rope=axes_dims_rope,
             theta=theta,
+            quant_config=quant_config,
         )
 
     def forward(

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -130,6 +130,8 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
         ).to(self._execution_device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxKontextTransformer2DModel)
+        transformer_kwargs["od_config"] = od_config
+        transformer_kwargs["quant_config"] = od_config.quantization_config
         self.transformer = FluxKontextTransformer2DModel(**transformer_kwargs)
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -5,7 +5,7 @@
 import math
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
@@ -29,6 +29,11 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = init_logger(__name__)
 
 
@@ -96,7 +101,15 @@ class DistributedRMSNorm(nn.Module):
 class ColumnParallelGELU(nn.Module):
     """Column parallel linear with GELU activation."""
 
-    def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int,
+        *,
+        approximate: str = "tanh",
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
         super().__init__()
         self.proj = ColumnParallelLinear(
             dim_in,
@@ -104,6 +117,7 @@ class ColumnParallelGELU(nn.Module):
             bias=bias,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.approximate = approximate
 
@@ -121,18 +135,15 @@ class HeliosFeedForward(nn.Module):
         inner_dim: int,
         dim_out: int | None = None,
         bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         dim_out = dim_out or dim
 
-        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
+        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config)
         self.net_1 = nn.Identity()
         self.net_2 = RowParallelLinear(
-            inner_dim,
-            dim_out,
-            bias=bias,
-            input_is_parallel=True,
-            return_bias=False,
+            inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -262,6 +273,7 @@ class HeliosSelfAttention(nn.Module):
         dropout: float = 0.0,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -275,6 +287,7 @@ class HeliosSelfAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             bias=True,
+            quant_config=quant_config,
         )
 
         self.num_heads = self.to_qkv.num_heads
@@ -290,6 +303,7 @@ class HeliosSelfAttention(nn.Module):
             bias=True,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.dropout = nn.Dropout(dropout)
 
@@ -368,6 +382,7 @@ class HeliosCrossAttention(nn.Module):
         head_dim: int,
         eps: float = 1e-5,
         dropout: float = 0.0,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -382,6 +397,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.to_k = ColumnParallelLinear(
             dim,
@@ -389,6 +405,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.to_v = ColumnParallelLinear(
             dim,
@@ -396,6 +413,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
 
         tp_size = get_tensor_model_parallel_world_size()
@@ -411,6 +429,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.dropout = nn.Dropout(dropout)
 
@@ -461,6 +480,7 @@ class HeliosTransformerBlock(nn.Module):
         guidance_cross_attn: bool = False,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -475,19 +495,17 @@ class HeliosTransformerBlock(nn.Module):
             eps=eps,
             is_amplify_history=is_amplify_history,
             history_scale_mode=history_scale_mode,
+            quant_config=quant_config,
         )
 
         # 2. Cross-attention
         self.attn2 = HeliosCrossAttention(
-            dim=dim,
-            num_heads=num_heads,
-            head_dim=head_dim,
-            eps=eps,
+            dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config
         )
         self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 3. Feed-forward
-        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
+        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config)
         self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -602,6 +620,7 @@ class HeliosTransformer3DModel(nn.Module):
         has_multi_term_memory_patch: bool = True,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
 
@@ -689,6 +708,7 @@ class HeliosTransformer3DModel(nn.Module):
                     guidance_cross_attn=guidance_cross_attn,
                     is_amplify_history=is_amplify_history,
                     history_scale_mode=history_scale_mode,
+                    quant_config=quant_config,
                 )
                 for _ in range(num_layers)
             ]

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -9,7 +9,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -31,6 +31,11 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -80,7 +85,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
     return load_json_config(model_path, subfolder, "config.json", local_files_only)
 
 
-def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
+def create_transformer_from_config(
+    config: dict, quant_config: QuantizationConfig | None = None
+) -> HeliosTransformer3DModel:
     kwargs = {}
 
     key_map = [
@@ -112,7 +119,7 @@ def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
                 val = tuple(val)
             kwargs[key] = val
 
-    return HeliosTransformer3DModel(**kwargs)
+    return HeliosTransformer3DModel(quant_config=quant_config, **kwargs)
 
 
 def get_helios_post_process_func(
@@ -192,7 +199,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
         ).to(self.device)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        self.transformer = create_transformer_from_config(
+            transformer_config, quant_config=od_config.quantization_config
+        )
 
         # Read scheduler config to determine scheduler type
         sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only)

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py

@@ -3,12 +3,28 @@
 
 from typing import Any
 
+import vllm.forward_context as _vllm_fc
 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.platforms import current_omni_platform
 
 
+def _set_forward_context_num_tokens(num_tokens: int) -> None:
+    """Set num_tokens on the vLLM ForwardContext for MoE routing.
+
+    After the rebase to vLLM 0.18.0, SharedFusedMoE expects
+    ForwardContext.num_tokens to be set. Without it, MoE expert
+    routing may produce incorrect results (silent correctness bug).
+    """
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    forward_context.num_tokens = num_tokens
+    if not hasattr(forward_context, "in_profile_run"):
+        forward_context.in_profile_run = False
+
+
 class HunyuanFusedMoEDefault(SharedFusedMoE):
     def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
         super().__init__(prefix=prefix, **kwargs)
@@ -21,6 +37,7 @@ class HunyuanFusedMoEDefault(SharedFusedMoE):
         self._init_hook_handle.remove()
 
     def forward(self, hidden_states: Any, router_logits: Any) -> Any:
+        _set_forward_context_num_tokens(hidden_states.shape[0])
         return super().forward(hidden_states, router_logits)
 
 

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -2667,6 +2667,15 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
 
+        # ---- TeaCache initialization ----
+        tea_cache_config = getattr(self.model, "_tea_cache_config", None)
+        if tea_cache_config is not None:
+            tc_rescale = np.poly1d(tea_cache_config.coefficients)
+            tc_acc_dist = 0.0
+            tc_prev_mod = None
+            tc_prev_pred = None
+            tc_cnt = 0
+
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 # expand the latents if we are doing classifier free guidance
@@ -2675,17 +2684,42 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
 
                 t_expand = t.repeat(latent_model_input.shape[0])
 
-                model_inputs = self.model.prepare_inputs_for_generation(
-                    input_ids,
-                    images=latent_model_input,
-                    timestep=t_expand,
-                    **model_kwargs,
-                )
+                # ---- TeaCache: decide whether to compute or reuse ----
+                should_compute = True
+                if tea_cache_config is not None:
+                    with torch.no_grad():
+                        # Use timestep embedding as the modulated input for cache decision
+                        cur_mod = self.model.time_embed(t.unsqueeze(0))
+                    if tc_cnt > 0 and tc_prev_mod is not None:
+                        rel_dist = (
+                            ((cur_mod - tc_prev_mod).abs().mean() / (tc_prev_mod.abs().mean() + 1e-8)).cpu().item()
+                        )
+                        rescaled = float(tc_rescale(rel_dist))
+                        tc_acc_dist += abs(rescaled)
+                        if tc_acc_dist < tea_cache_config.rel_l1_thresh:
+                            should_compute = False
+                        else:
+                            tc_acc_dist = 0.0
+                    tc_prev_mod = cur_mod.detach()
 
-                with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-                    model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
-                    pred = model_output["diffusion_prediction"]
-                pred = pred.to(dtype=torch.float32)
+                if should_compute:
+                    model_inputs = self.model.prepare_inputs_for_generation(
+                        input_ids,
+                        images=latent_model_input,
+                        timestep=t_expand,
+                        **model_kwargs,
+                    )
+
+                    with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+                        model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
+                        pred = model_output["diffusion_prediction"]
+                    pred = pred.to(dtype=torch.float32)
+
+                    if tea_cache_config is not None:
+                        tc_prev_pred = pred.clone()
+                else:
+                    # TeaCache fast path: reuse previous prediction
+                    pred = tc_prev_pred
 
                 # perform guidance
                 if self.do_classifier_free_guidance:
@@ -2695,7 +2729,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 # compute the previous noisy sample x_t -> x_t-1
                 latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]
 
-                if i != len(timesteps) - 1:
+                if i != len(timesteps) - 1 and should_compute:
                     model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                         model_output,
                         model_kwargs,
@@ -2709,6 +2743,9 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     model_kwargs["query_lens"] = query_lens
                     model_kwargs["seq_lens"] = seq_lens
 
+                if tea_cache_config is not None:
+                    tc_cnt += 1
+
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
                     for k in callback_on_step_end_tensor_inputs:

vllm_omni/entrypoints/openai/api_server.py

@@ -1207,38 +1207,16 @@ _remove_route_from_router(router, "/v1/models")
 
 @router.get("/v1/models")
 async def show_available_models(raw_request: Request) -> JSONResponse:
-    """Show available models endpoint that works for both LLM and diffusion modes.
+    """Show available models for both LLM and diffusion modes.
 
-    Returns model information in OpenAI-compatible format.
+    Delegates to state.openai_serving_models which is set to either
+    OpenAIServingModels (LLM) or _DiffusionServingModels (pure diffusion).
     """
-    # Check if we're in diffusion mode
-    diffusion_model_name = getattr(raw_request.app.state, "diffusion_model_name", None)
-    if diffusion_model_name is not None:
-        # Diffusion mode - return the loaded model
-        return JSONResponse(
-            content={
-                "object": "list",
-                "data": [
-                    {
-                        "id": diffusion_model_name,
-                        "object": "model",
-                        "created": 0,
-                        "owned_by": "vllm-omni",
-                        "permission": [],
-                    }
-                ],
-            }
-        )
-
-    # LLM mode - delegate to openai_serving_models
-    openai_serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
-    if openai_serving_models is not None:
-        models = await openai_serving_models.show_available_models()
+    handler = getattr(raw_request.app.state, "openai_serving_models", None)
+    if handler is not None:
+        models = await handler.show_available_models()
         return JSONResponse(content=models.model_dump())
-
-    return JSONResponse(
-        content={"object": "list", "data": []},
-    )
+    return JSONResponse(content={"object": "list", "data": []})
 
 
 # Image generation API endpoints
@@ -1775,10 +1753,27 @@ def _choose_output_format(output_format: str | None, background: str | None) ->
     return "jpeg"
 
 
+def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image:
+    fmt = format.lower()
+    if fmt not in {"jpg", "jpeg"}:
+        return image
+
+    if image.mode == "RGB":
+        return image
+
+    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
+        alpha_image = image.convert("RGBA")
+        flattened = Image.new("RGB", alpha_image.size, (255, 255, 255))
+        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
+        return flattened
+
+    return image.convert("RGB")
+
+
 def _encode_image_base64_with_compression(
     image: Image.Image, format: str = "png", output_compression: int = 100
 ) -> str:
-    """Encode PIL Image to base64 PNG string.
+    """Encode PIL Image to a base64 image string.
 
     Args:
         image: PIL Image object
@@ -1788,6 +1783,7 @@ def _encode_image_base64_with_compression(
         Base64-encoded image as string
     """
     buffer = io.BytesIO()
+    image = _prepare_image_for_output_format(image, format)
     save_kwargs = {}
     if format in ("jpg", "jpeg", "webp"):
         save_kwargs["quality"] = output_compression

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -1,9 +1,10 @@
 """Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache.
 
 * SDPA attention (F.scaled_dot_product_attention) with native GQA support
+* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers)
 * Per-call embedding buffer to avoid cross-request aliasing
 * Pre-allocated position_ids (read-only, safe to persist)
-* torch.compile on inner transformer by default
+* torch.compile (epilogue_fusion=False) on inner transformer by default
 * Inline sampling (top-k + top-p) -- no custom op overhead
 """
 
@@ -12,15 +13,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    MergedColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.platforms import current_omni_platform
 
@@ -28,15 +22,84 @@ logger = init_logger(__name__)
 
 
 # ===================================================================
-#  Standalone Attention (SDPA, no KV cache, no HF backend fallback)
+# HF-numerics-compatible layers for code predictor
 # ===================================================================
+#
+# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32,
+# rotate_half RoPE) to produce outputs numerically identical to the
+# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel,
+# get_rope) introduce small precision differences that compound across
+# the autoregressive steps of the code predictor, causing severe
+# audio quality degradation.
+#
+# See: https://github.com/vllm-project/vllm-omni/issues/2274
+
+
+class _RMSNorm(nn.Module):
+    """RMSNorm matching HuggingFace's implementation exactly.
+
+    Computes variance in float32 to avoid bfloat16 precision loss.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class _RotaryEmbedding(nn.Module):
+    """RoPE matching HuggingFace's implementation exactly.
+
+    Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False).
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__()
+        head_dim = getattr(
+            config,
+            "head_dim",
+            config.hidden_size // config.num_attention_heads,
+        )
+        rope_theta = getattr(config, "rope_theta", 10000.0)
+        inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # position_ids: [batch, seq_len]
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        # Force float32 (matching HF)
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
 class Qwen3OmniCodePredictorAttention(nn.Module):
     """Multi-head self-attention for code predictor.
 
-    Uses ``F.scaled_dot_product_attention`` directly.  No KV cache -- the code
-    predictor always re-prefills the full (short) sequence each AR step.
+    Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm.
+    No KV cache -- the code predictor always re-prefills the full (short)
+    sequence each AR step.
 
     Input : [B, seq_len, hidden_size]
     Output: [B, seq_len, hidden_size]
@@ -45,7 +108,6 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
     def __init__(
         self,
         config,
-        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -58,58 +120,52 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
             cp_cfg.hidden_size // cp_cfg.num_attention_heads,
         )
         self.hidden_size = cp_cfg.hidden_size
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self._use_gqa = self.num_kv_heads != self.num_heads
 
-        self.qkv_proj = QKVParallelLinear(
-            hidden_size=self.hidden_size,
-            head_size=self.head_dim,
-            total_num_heads=self.num_heads,
-            total_num_kv_heads=self.num_kv_heads,
+        # Separate q/k/v projections matching HF (no fused packing)
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim,
             bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
-            disable_tp=True,
         )
-        self.o_proj = RowParallelLinear(
-            input_size=self.num_heads * self.head_dim,
-            output_size=self.hidden_size,
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
             bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-            disable_tp=True,
         )
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            max_position=cp_cfg.max_position_embeddings,
-            rope_parameters=None,
-            dual_chunk_attention_config=None,
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
         )
-        self.q_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+        )
+        self.q_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
+        self.k_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
         bsz, seq_len, _ = hidden_states.shape
+        hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
+        hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)
 
-        # Flatten to 2-D so vLLM rotary_emb gets [num_tokens, size]
-        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
+        k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)
 
-        # QK-norm -> RoPE (both 2-D)
-        q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
-        k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
-        q, k = self.rotary_emb(position_ids, q, k)
-
-        # [B, heads, seq, head_dim]
-        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        cos, sin = position_embeddings
+        # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
+        cos = cos.unsqueeze(1)  # [batch, 1, seq_len, head_dim]
+        sin = sin.unsqueeze(1)
+        q = (q * cos) + (_rotate_half(q) * sin)
+        k = (k * cos) + (_rotate_half(k) * sin)
 
         attn_out = F.scaled_dot_product_attention(
             q,
@@ -120,9 +176,9 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
             enable_gqa=self._use_gqa,
         )
 
-        attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
-        output, _ = self.o_proj(attn_out)
-        return output.view(bsz, seq_len, -1)
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.o_proj(attn_out)
+        return output
 
 
 # ===================================================================
@@ -131,40 +187,23 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
 
 
 class Qwen3OmniCodePredictorMLP(nn.Module):
-    """SiLU-gated MLP for code predictor."""
+    """SiLU-gated MLP for code predictor, matching HF's implementation."""
 
     def __init__(
         self,
         config,
-        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
         hidden_size = config.code_predictor_config.hidden_size
         intermediate_size = config.code_predictor_config.intermediate_size
 
-        self.gate_up_proj = MergedColumnParallelLinear(
-            input_size=hidden_size,
-            output_sizes=[intermediate_size, intermediate_size],
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-            disable_tp=True,
-        )
-        self.down_proj = RowParallelLinear(
-            input_size=intermediate_size,
-            output_size=hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-            disable_tp=True,
-        )
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        gate_up, _ = self.gate_up_proj(hidden_states)
-        gate, up = gate_up.chunk(2, dim=-1)
-        down, _ = self.down_proj(F.silu(gate) * up)
-        return down
+        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
 
 
 # ===================================================================
@@ -178,32 +217,29 @@ class Qwen3OmniCodePredictorDecoderLayer(nn.Module):
     def __init__(
         self,
         config,
-        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
         self.self_attn = Qwen3OmniCodePredictorAttention(
             config,
-            quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
         self.mlp = Qwen3OmniCodePredictorMLP(
             config,
-            quant_config=quant_config,
             prefix=f"{prefix}.mlp",
         )
         cp_cfg = config.code_predictor_config
-        self.input_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
+        self.input_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
+        self.post_attention_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, position_ids)
+        hidden_states = self.self_attn(hidden_states, position_embeddings)
         hidden_states = residual + hidden_states
 
         residual = hidden_states
@@ -238,13 +274,13 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
             [
                 Qwen3OmniCodePredictorDecoderLayer(
                     vllm_config.model_config.hf_config,
-                    quant_config=vllm_config.quant_config,
                     prefix=f"{prefix}.layers.{idx}",
                 )
                 for idx in range(config.num_hidden_layers)
             ]
         )
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = _RotaryEmbedding(config)
 
     def forward(
         self,
@@ -252,8 +288,9 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         position_ids: torch.Tensor,
     ) -> torch.Tensor:
         hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
         for layer in self.layers:
-            hidden_states = layer(hidden_states, position_ids)
+            hidden_states = layer(hidden_states, position_embeddings)
         hidden_states = self.norm(hidden_states)
         return hidden_states
 
@@ -332,7 +369,8 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
         if self._pos_ids is not None and self._pos_ids.device == device:
             return
         max_seq = self.num_code_groups + 1
-        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device)
+        # [1, max_seq] for HF-style RoPE (will be expanded to [bsz, seq_len] at use)
+        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0)
 
     def _ensure_cached_refs(self) -> None:
         if self._lm_heads is not None:
@@ -344,16 +382,18 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
         if self._model_fwd is not None:
             return
         if current_omni_platform.supports_torch_inductor():
-            # torch.compile fuses the 5-layer transformer's small kernels,
-            # reducing BF16 intermediate round-trips and improving precision.
-            # mode="default" avoids Inductor's own CUDA graph capture so it
-            # doesn't conflict with the outer CUDAGraphWrapper.
+            # torch.compile fuses RMSNorm/RoPE in ways that lose float32
+            # precision, compounding across AR steps. Use epilogue_fusion=False
+            # to disable the problematic fusions while still getting kernel
+            # fusion benefits for the linear layers and SDPA.
             self._model_fwd = torch.compile(
                 self.model.forward,
-                mode="default",
                 dynamic=True,
+                options={
+                    "epilogue_fusion": False,
+                },
             )
-            logger.info("code_predictor: torch.compile enabled (mode=default)")
+            logger.info("code_predictor: torch.compile enabled (no epilogue fusion)")
         else:
             self._model_fwd = self.model.forward
             logger.info("code_predictor: using eager mode (no torch.compile)")
@@ -413,8 +453,8 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
         for step in range(1, num_groups):
             seq_len = step + 1
             projected = proj_buf[:bsz, :seq_len, :]
-            step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
-            assert step_pos_ids.shape[0] == bsz * seq_len
+            # position_ids: [batch, seq_len] for HF-style RoPE
+            step_pos_ids = pos_ids[:, :seq_len].expand(bsz, -1)
 
             hidden_out = model_fwd(projected, step_pos_ids)
 
@@ -446,21 +486,11 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
     # ------------------------------------------------------------------
 
     def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights with mapping for fused QKV and gate_up projections.
+        """Load weights directly (no fused projection remapping needed).
 
-        Maps original HF weights (q_proj, k_proj, v_proj, gate_proj, up_proj)
-        to fused vLLM weights (qkv_proj, gate_up_proj).
+        Since we use separate nn.Linear for q/k/v/o and gate/up/down,
+        weight names match the HF checkpoint directly.
         """
-        # Mapping for fused projections
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
@@ -469,36 +499,12 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
             if "rotary_emb.inv_freq" in name:
                 continue
 
-            # Handle stacked/fused parameters
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                # Skip if parameter doesn't exist (e.g., bias)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Non-stacked parameters - use default loading
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", None)
-                if weight_loader is not None:
-                    weight_loader(param, loaded_weight)
-                else:
-                    param.data.copy_(loaded_weight)
+            param = params_dict.get(name)
+            if param is None:
+                continue
 
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
             loaded_params.add(name)
 
         return loaded_params

vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py

@@ -482,8 +482,6 @@ class Qwen3TTSConfig(PretrainedConfig):
         tts_eos_token_id=151673,
         **kwargs,
     ):
-        super().__init__(**kwargs)
-
         if talker_config is None:
             talker_config = {}
             logger.info("talker_config is None. Initializing talker model with default values")
@@ -494,6 +492,8 @@ class Qwen3TTSConfig(PretrainedConfig):
         self.talker_config = Qwen3TTSTalkerConfig(**talker_config)
         self.speaker_encoder_config = Qwen3TTSSpeakerEncoderConfig(**speaker_encoder_config)
 
+        super().__init__(**kwargs)
+
         self.tokenizer_type = tokenizer_type
         self.tts_model_size = tts_model_size
         self.tts_model_type = tts_model_type