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

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

.buildkite/pipeline-intel.yaml

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

.buildkite/test-amd-merge.yml

@@ -7,7 +7,9 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
 
 - label: "Diffusion Model Test"
   agent_pool: mi325_2

.buildkite/test-amd-ready.yaml

@@ -7,7 +7,9 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+    # that is not supported on CPU
+    - "timeout 20m pytest -vvvv -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
 
 - label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1

.github/mergify.yml

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

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.20.0-cu130
+ARG VLLM_BASE_TAG=v0.20.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.cuda

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

docker/Dockerfile.rocm

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

docker/Dockerfile.xpu

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

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.20.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm700
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```
 
 #### Installation of vLLM-Omni
@@ -37,10 +37,10 @@ uv pip install onnxruntime-rocm
 If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
 
 ```bash
-uv pip install vllm==0.20.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm700
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```
 
-The pre-built 0.20.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source

docs/getting_started/quickstart.md

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

docs/user_guide/quantization/fp8.md

@@ -32,15 +32,15 @@ guide. FP8 on Ampere may use a weight-only path where available.
 
 ### Diffusion Model (Qwen-Image, Wan2.2)
 
-| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
-|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
-| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
-| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
-| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
-| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
-| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
-| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
-| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
+| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` |
+|-------|-----------|:-------:|:------:|----------------|------------------|
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None |
+| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None |
+| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None |
+| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

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

examples/offline_inference/helios/end2end.py

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

examples/offline_inference/image_to_image/image_edit.py

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

examples/offline_inference/image_to_video/image_to_video.py

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

examples/offline_inference/magi_human/end2end.py

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

examples/offline_inference/text_to_image/text_to_image.py

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

examples/offline_inference/text_to_video/text_to_video.py

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

examples/offline_inference/vace/vace_video_generation.py

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

pyproject.toml

@@ -32,7 +32,6 @@ classifiers = [
 # Dependencies are now managed dynamically via setup.py based on detected hardware platform.
 # This allows automatic installation of the correct platform-specific dependencies (CUDA/ROCm/CPU/XPU/NPU)
 # without requiring extras like [cuda]. See requirements/ directory for platform-specific dependencies.
-# Note: vllm is intentionally excluded due to entrypoints overwrite issue.
 
 [project.optional-dependencies]
 
@@ -90,7 +89,6 @@ Documentation = "https://vllm-omni.readthedocs.io"
 "Bug Tracker" = "https://github.com/vllm-project/vllm-omni/issues"
 
 [project.scripts]
-vllm = "vllm_omni.entrypoints.cli.main:main"
 vllm-omni = "vllm_omni.entrypoints.cli.main:main"
 
 

tests/diffusion/models/bagel/test_trajectory_recording.py

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

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -67,7 +67,7 @@ if current_omni_platform.is_rocm():
     ]
 
 # Maximum allowed difference per color channel
-PIXEL_TOLERANCE = 5
+PIXEL_TOLERANCE = 10
 
 # Default test prompt
 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"

tests/e2e/online_serving/test_mimo_audio.py

@@ -83,6 +83,7 @@ def get_max_batch_size(size_type="few"):
 @pytest.mark.omni
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.skip(reason="CI failed 8571")
 def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Test audio and text input processing and text/audio output generation via OpenAI API.

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -356,7 +356,6 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.skip(reason="issue: #2827")
 def test_audio_in_video_001(omni_server, openai_client) -> None:
     """
     Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
@@ -376,7 +375,7 @@ def test_audio_in_video_001(omni_server, openai_client) -> None:
         "messages": messages,
         "stream": False,
         "use_audio_in_video": True,
-        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
+        "key_words": {"video": VIDEO_KEY},
     }
     openai_client.send_omni_request(request_config)
 

tests/e2e/online_serving/test_zimage_expansion.py

@@ -98,7 +98,7 @@ def _get_diffusion_feature_cases():
                 ],
             ),
             id="layerwise_hsdp",
-            marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
+            marks=[*FOUR_CARD_MARKS],
         ),
     ]
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -55,6 +55,44 @@ def test_default_stage_devices_from_sequence_parallel():
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
 
 
+def test_default_stage_config_uses_parallel_size_kwargs():
+    """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "pipeline_parallel_size": 2,
+            "data_parallel_size": 3,
+            "tensor_parallel_size": 4,
+            "enable_expert_parallel": True,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.pipeline_parallel_size == 2
+    assert parallel_config.data_parallel_size == 3
+    assert parallel_config.tensor_parallel_size == 4
+    assert parallel_config.enable_expert_parallel is True
+
+
+def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
+    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "pipeline_parallel_size": None,
+            "data_parallel_size": None,
+            "tensor_parallel_size": None,
+            "enable_expert_parallel": None,
+            "enforce_eager": None,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.pipeline_parallel_size == 1
+    assert parallel_config.data_parallel_size == 1
+    assert parallel_config.tensor_parallel_size == 1
+    assert parallel_config.enable_expert_parallel is False
+    assert stage_cfg["engine_args"]["enforce_eager"] is False
+
+
 def test_default_stage_config_propagates_ulysses_mode():
     """Ensure UAA mode survives default diffusion-stage creation."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(

tests/examples/offline_inference/test_text_to_image.py

@@ -26,27 +26,9 @@ def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool,
 
 README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
 
-# Track: https://github.com/vllm-project/vllm-omni/issues/3123 — CLI example subprocess exit 1 in CI; re-enable when fixed.
-_SKIP_T2I_SNIPPET_IDS = frozenset(
-    {
-        "quick_start_002",
-        "more_cli_examples_001",
-        "more_cli_examples_002",
-        "more_cli_examples_003",
-        "more_cli_examples_007",
-        "more_cli_examples_008",
-    }
-)
-_SKIP_T2I_REASON = (
-    "Skipped: text_to_image README examples failing in CI (subprocess exit 1). "
-    "See https://github.com/vllm-project/vllm-omni/issues/3123"
-)
-
 
 @pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
 def test_text_to_image(snippet: ReadmeSnippet, example_runner: ExampleRunner):
-    if snippet.test_id in _SKIP_T2I_SNIPPET_IDS:
-        pytest.skip(_SKIP_T2I_REASON)
     should_skip, reason = snippet.skip
     if should_skip:
         pytest.skip(reason)

tests/helpers/mark.py

@@ -3,14 +3,7 @@
 """Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …)."""
 
 import pytest
-
-try:
-    from vllm.utils.torch_utils import cuda_device_count_stateless
-except ImportError:
-    import torch
-
-    def cuda_device_count_stateless() -> int:
-        return torch.cuda.device_count()
+from vllm.platforms import current_platform
 
 # Re-exported from tests.helpers.env (GPU wait + DeviceMemoryMonitor).
 
@@ -27,8 +20,9 @@ def cuda_marks(*, res: str, num_cards: int):
     if num_cards == 1:
         return marks
     test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards)
+
     test_skipif = pytest.mark.skipif_cuda(
-        cuda_device_count_stateless() < num_cards,
+        not current_platform.is_cuda() or (current_platform.device_count() < num_cards),
         reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
     )
     return marks + [test_distributed, test_skipif]

vllm_omni/config/stage_config.py

@@ -1009,6 +1009,8 @@ class StageConfigFactory:
             cli_overrides = {}
 
         trust_remote_code = cli_overrides.get("trust_remote_code", True)
+        if trust_remote_code is None:
+            trust_remote_code = False
 
         # --- New path: check pipeline registry by model_type first ---
         model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
@@ -1352,8 +1354,8 @@ class StageConfigFactory:
 
             hf_config = get_config(model, trust_remote_code=trust_remote_code)
             return hf_config.model_type, hf_config
-        except Exception:
-            pass
+        except Exception as e:
+            logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path")
 
         # Fallback: read config.json directly for custom model types that
         # are not registered with transformers (e.g. qwen3_tts).
@@ -1390,8 +1392,8 @@ class StageConfigFactory:
                             class_name,
                         )
                         return pipeline_cfg.model_type, None
-        except Exception:
-            pass
+        except Exception as e:
+            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
 
         # Final fallback: some models (e.g. CosyVoice3) ship an empty
         # config.json and rely on naming conventions. Match the model path

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py

@@ -28,12 +28,6 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
         model.init_distributed()
         return model
 
-    @classmethod
-    def from_config(cls, *args: Any, **kwargs: Any):
-        model = super().from_config(*args, **kwargs)
-        model.init_distributed()
-        return model
-
     def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
         # mostly copy from AutoencoderKL
         overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -385,9 +385,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise
 
     def check_health(self) -> None:
-        self._ensure_open()
         if self.is_failed:
             raise EngineDeadError()
+        self._ensure_open()
         for p in self._processes:
             if not p.is_alive():
                 self.is_failed = True

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 from PIL import Image
 from vllm.logger import init_logger
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
@@ -329,6 +330,12 @@ class InlineStageDiffusionClient:
             None,
         )
 
+    def check_health(self) -> None:
+        """Check if the inline diffusion engine and its workers are healthy."""
+        if self._shutting_down:
+            raise EngineDeadError("InlineStageDiffusionClient is shutting down")
+        self._engine.executor.check_health()
+
     def shutdown(self) -> None:
         self._shutting_down = True
 

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -26,7 +26,6 @@ from vllm.model_executor.model_loader.weight_utils import (
     multi_thread_safetensors_weights_iterator,
     safetensors_weights_iterator,
 )
-from vllm.transformers_utils.repo_utils import file_exists
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import set_default_torch_dtype
 
@@ -50,8 +49,6 @@ def _natural_sort_key(filepath: str) -> list:
 
 MODEL_INDEX = "model_index.json"
 DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
-TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
-INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
 
 
 class DiffusersPipelineLoader:
@@ -102,22 +99,8 @@ class DiffusersPipelineLoader:
         is_local = os.path.isdir(model_name_or_path)
         load_format = self.load_config.load_format
         use_safetensors = False
-        possible_index_files = [
-            f"{subfolder}/{index_file}" if subfolder is not None else index_file for index_file in INDEX_FILES
-        ]
-        available_index_file = list(
-            filter(lambda f: file_exists(model_name_or_path, f, revision=revision), possible_index_files)
-        )
-        if len(available_index_file) > 1:
-            raise ValueError(
-                f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: {available_index_file}"
-            )
-        index_file_with_subfolder = available_index_file[0] if len(available_index_file) == 1 else None
-        index_file = (
-            index_file_with_subfolder.split("/")[-1]
-            if index_file_with_subfolder and subfolder is not None
-            else index_file_with_subfolder
-        )
+        index_file = DIFFUSION_MODEL_WEIGHTS_INDEX
+        index_file_with_subfolder = f"{subfolder}/{index_file}" if subfolder else index_file
 
         # only hf is supported currently
         if load_format == "auto":
@@ -180,7 +163,7 @@ class DiffusersPipelineLoader:
             hf_weights_files = filter_duplicate_safetensors_files(
                 hf_weights_files,
                 filter_folder,
-                index_file or "",
+                index_file,
             )
         else:
             hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1842,6 +1842,8 @@ class Bagel(nn.Module):
         # ── SP + CFG: sequential single-branch forwards ──
         use_sp = self._sp_size > 1
         if use_sp and use_cfg_text:
+            if return_trajectory_latents and len(timesteps) > 0:
+                trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
                 in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1]
@@ -1905,13 +1907,15 @@ class Bagel(nn.Module):
                     x_t = x_t - v_t.to(x_t.device) * dts[i]
                 if return_trajectory_latents:
                     trajectory_latents.append(x_t.clone())
-                    trajectory_timesteps.append(timesteps[i] - dts[i])
+                    trajectory_timesteps.append(timesteps[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
             return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
 
         # ── SP without CFG: direct single-branch loop ──
         if use_sp:
+            if return_trajectory_latents and len(timesteps) > 0:
+                trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
                 v_t = self.forward_single_branch(
@@ -1938,7 +1942,7 @@ class Bagel(nn.Module):
                     x_t = x_t - v_t.to(x_t.device) * dts[i]
                 if return_trajectory_latents:
                     trajectory_latents.append(x_t.clone())
-                    trajectory_timesteps.append(timesteps[i] - dts[i])
+                    trajectory_timesteps.append(timesteps[i])
 
             unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
             return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
@@ -1999,6 +2003,9 @@ class Bagel(nn.Module):
                 "merged_cache": self._merge_naive_caches(branches_cache),
             }
 
+        if return_trajectory_latents and len(timesteps) > 0:
+            trajectory_latents.append(x_t.clone())
+
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
             if t > cfg_interval[0] and t <= cfg_interval[1]:
@@ -2036,7 +2043,7 @@ class Bagel(nn.Module):
                 x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
-                trajectory_timesteps.append(timesteps[i] - dts[i])
+                trajectory_timesteps.append(timesteps[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
         return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
@@ -2137,6 +2144,9 @@ class Bagel(nn.Module):
         )
         _sched_kw = scheduler_kwargs or {}
 
+        if return_trajectory_latents and len(timesteps) > 0:
+            trajectory_latents.append(x_t.clone())
+
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
             use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
@@ -2194,7 +2204,7 @@ class Bagel(nn.Module):
                 x_t = x_t - v_t.to(x_t.device) * dts[i]
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
-                trajectory_timesteps.append(timesteps[i] - dts[i])
+                trajectory_timesteps.append(timesteps[i])
 
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
         return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs

vllm_omni/diffusion/models/utils.py

@@ -5,132 +5,6 @@ from __future__ import annotations
 
 import json
 import os
-from typing import TYPE_CHECKING, Literal
-
-import torch
-import torch.nn as nn
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    ReplicatedLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.models.transformers.utils import init_on_device_without_buffers
-from vllm.model_executor.models.utils import maybe_prefix
-
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
-
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedModel
-    from transformers.models.auto.auto_factory import _BaseAutoModelClass
-    from vllm.model_executor.layers.quantization.base_config import (
-        QuantizationConfig,
-    )
-
-
-Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]
-
-
-def replace_linear_class(
-    linear: nn.Linear,
-    style: Style = "replicate",
-    quant_config: QuantizationConfig | None = None,
-    *,
-    prefix: str = "",
-) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
-    """
-    Replace nn.Linear with one of vLLM's tensor parallel linear classes.
-
-    Args:
-        linear: `nn.Linear` to be replaced.
-        style: Tensor parallel style of the new linear, e.g. "colwise".
-        quant_config: Quantization config for the new linear.
-    Returns:
-        The new linear.
-    """
-
-    if not isinstance(style, str):
-        raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")
-
-    vllm_linear_maps = {
-        "colwise": (ColumnParallelLinear, {}),
-        "colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
-        "rowwise": (RowParallelLinear, {}),
-        "rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
-        "replicate": (ReplicatedLinear, {}),
-    }
-    vllm_linear_cls, vllm_linear_kwargs = vllm_linear_maps[style]
-
-    return vllm_linear_cls(
-        input_size=linear.in_features,
-        output_size=linear.out_features,
-        bias=linear.bias is not None,
-        quant_config=quant_config,
-        prefix=prefix,
-        return_bias=False,
-        **vllm_linear_kwargs,
-    )
-
-
-def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig):
-    """Recursively replace modules in the model as needed.
-    Currently, this replaces:
-    - `nn.Linear` with vLLM's tensor parallel linear classes
-    """
-    # Prefix the patterns because we always start from `self.model`
-    quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
-
-    def _recursive_replace(module: nn.Module, prefix: str):
-        for child_name, child_module in module.named_children():
-            new_module = child_module
-            qual_name = maybe_prefix(prefix, child_name)
-            # Replace modules as needed
-            if isinstance(child_module, nn.Linear):
-                style = "replicate"
-                new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name)
-            else:
-                _recursive_replace(child_module, prefix=qual_name)
-            if new_module is not child_module:
-                setattr(module, child_name, new_module)
-
-    _recursive_replace(model, prefix="")
-
-
-def init_parameters(
-    module: nn.Module,
-    dtype: torch.dtype | None,
-    device: torch.device | None = None,
-):
-    for name, param in module.named_parameters(recurse=False):
-        if param.device == torch.device("meta"):
-            new_param = nn.Parameter(
-                torch.empty_like(
-                    param.data,
-                    dtype=dtype,
-                    device=device,
-                ),
-                requires_grad=param.requires_grad,
-            )
-            setattr(module, name, new_param)
-    for child in module.children():
-        init_parameters(child, dtype, device)
-
-
-def create_transformers_model(
-    auto_cls: _BaseAutoModelClass,
-    od_config: OmniDiffusionConfig,
-    hf_config: PretrainedConfig,
-    dtype: torch.dtype | None = None,
-    device: torch.device | None = None,
-) -> PreTrainedModel:
-    """Create a HuggingFace model using the given auto class and model name."""
-    dtype = dtype or od_config.dtype
-    device = device or torch.get_default_device()
-    with init_on_device_without_buffers("meta"):
-        model = auto_cls.from_config(hf_config)
-    recursive_replace_linear(model, od_config)
-    init_parameters(model, dtype=dtype, device=device)
-    return model
 
 
 def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -28,7 +28,7 @@ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -36,7 +36,6 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import Distribu
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
-from vllm_omni.diffusion.models.utils import create_transformers_model
 from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
@@ -170,25 +169,13 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         super().__init__()
         self.od_config = od_config
         self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="text_encoder",
-                revision=od_config.revision,
-                prefix="text_encoder.",
-            ),
             DiffusersPipelineLoader.ComponentSource(
                 model_or_path=od_config.model,
                 subfolder="transformer",
-                revision=od_config.revision,
+                revision=None,
                 prefix="transformer.",
                 fall_back_to_pt=True,
-            ),
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="vae",
-                revision=od_config.revision,
-                prefix="vae.",
-            ),
+            )
         ]
         self._execution_device = get_local_device()
         model = od_config.model
@@ -205,19 +192,12 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
             model, subfolder="scheduler", local_files_only=local_files_only
         )
 
-        text_encoder_config = AutoConfig.from_pretrained(
+        self.text_encoder = AutoModel.from_pretrained(
             model, subfolder="text_encoder", local_files_only=local_files_only
-        )
-        self.text_encoder = create_transformers_model(
-            AutoModelForCausalLM,
-            od_config,
-            hf_config=text_encoder_config,
         ).to(self._execution_device)
-        if text_encoder_config.tie_word_embeddings:
-            self.text_encoder.lm_head.weight = self.text_encoder.get_input_embeddings().weight
-
-        vae_config = DistributedAutoencoderKL.load_config(model, subfolder="vae", local_files_only=local_files_only)
-        self.vae = DistributedAutoencoderKL.from_config(vae_config).to(self._execution_device)
+        self.vae = DistributedAutoencoderKL.from_pretrained(
+            model, subfolder="vae", local_files_only=local_files_only
+        ).to(self._execution_device)
         self.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 

vllm_omni/engine/async_omni_engine.py

@@ -1294,9 +1294,12 @@ class AsyncOmniEngine:
             ring_degree = normalized_kwargs.get("ring_degree") or 1
             ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
             sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
+            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
+            data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1
             tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
             vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
+            enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", False)
             hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
             hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1)
@@ -1304,9 +1307,10 @@ class AsyncOmniEngine:
                 sequence_parallel_size = ulysses_degree * ring_degree
 
             parallel_config = DiffusionParallelConfig(
-                pipeline_parallel_size=1,
-                data_parallel_size=1,
+                pipeline_parallel_size=pipeline_parallel_size,
+                data_parallel_size=data_parallel_size,
                 tensor_parallel_size=tensor_parallel_size,
+                enable_expert_parallel=enable_expert_parallel,
                 sequence_parallel_size=sequence_parallel_size,
                 ulysses_degree=ulysses_degree,
                 ring_degree=ring_degree,
@@ -1333,7 +1337,7 @@ class AsyncOmniEngine:
             "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
             "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
-            "enforce_eager": kwargs.get("enforce_eager", False),
+            "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),
             "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),

vllm_omni/entrypoints/openai/api_server.py

@@ -65,7 +65,7 @@ from vllm.entrypoints.openai.speech_to_text.serving import (
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
-from vllm.entrypoints.pooling.pooling.serving import ServingPooling as OpenAIServingPooling
+from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
 from vllm.entrypoints.serve.disagg.serving import ServingTokens
 
@@ -800,10 +800,9 @@ async def omni_init_app_state(
         else None
     )
     state.openai_serving_pooling = (
-        OpenAIServingPooling(
+        ServingPooling(
             engine_client,
             state.openai_serving_models,
-            state.openai_serving_render,
             supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
             chat_template=resolved_chat_template,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -489,7 +489,7 @@ class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
     """
     Qwen3 Omni MoE Talker language model.
 
-    Extends Qwen3MoeLLMForCausalLM (which already uses SharedFusedMoE with
+    Extends Qwen3MoeLLMForCausalLM (which already uses FusedMoE with
     shared-expert support) and replaces the text embedding / LM head with a
     codec embedding so the talker operates over audio-codec tokens instead
     of text tokens.

vllm_omni/platforms/npu/models/hunyuan_fused_moe.py

@@ -103,6 +103,9 @@ def prepare_hunyuan_fused_moe_runtime() -> None:
     _ensure_forward_context_attr("flash_comm_v1_enabled", bool, False)
 
 
+# NOTE: As of v0.20,0 vLLM has folded SharedFusedMoE -> FusedMoE and removed the class.
+# Since AscendSharedFusedMoE is a subclass of SharedFusedMoE, we should be careful
+# to ensure that this is updated correctly.
 class AscendHunyuanFusedMoE(AscendSharedFusedMoE):
     def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
         super().__init__(prefix=prefix, **kwargs)