[Misc] Keep serving changes out of example refactor

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

.buildkite/cuda/test-merge.yml

@@ -139,12 +139,7 @@ steps:
         timeout_in_minutes: 50
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-        # h100_2, not h100_1: advanced_model asserts transcribe every audio
-        # response with Whisper, and tests/helpers/media.py only uses an
-        # accelerator for Whisper when a spare GPU exists. On a single card it
-        # grades on CPU: ~8-10 min per clip vs ~30 s, which alone overruns the
-        # 50-minute job timeout (build 2956).
-        mirror_hardwares: h100_2
+        mirror_hardwares: h100_1
 
       - label: "Omni · MiniCPM-o 4.5 Duplex Test"
         source_file_dependencies:

.buildkite/cuda/test-ready.yml

@@ -128,21 +128,6 @@ steps:
               image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
               always-pull: true
               propagate-environment: true
-              # vLLM 0.27 added a pre-flight capacity check in
-              # distributed/device_communicators/shm_broadcast.py. Without this
-              # the container gets Docker's 64 MiB default and every test in
-              # this step dies at engine start with
-              #   RuntimeError: Insufficient space in /dev/shm:
-              #                 240 MiB required, 64 MiB available
-              # (9 failed in 8.7s on build #2949). vLLM 0.26 had no such check
-              # and allocated regardless, so this step passed on main while
-              # being under-provisioned the whole time — 0.27 turned a silent
-              # risk into an explicit error.
-              #
-              # This is the only inline docker plugin block in .buildkite/cuda/;
-              # every other step inherits common/ci_mirror_hardwares.yml, which
-              # already sets 8gb. Matching that rather than inventing a value.
-              shm-size: "8gb"
               environment:
                 - "HF_HOME=/fsx/hf_cache"
               volumes:

docker/Dockerfile.ci

@@ -1,24 +1,7 @@
-# Target vllm version: v0.27.0 (releases/v0.27.0 branch).
+# Target vllm version: v0.26.0 (releases/v0.26.0 branch).
 # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
-#
-# Release-boundary form: there is no precompiled-wheel pin here, because the
-# base image already ships exactly the target vLLM. Everything the runtime
-# needs — torch, flashinfer, nvJitLink, numpy — comes from that image as one
-# self-consistent set.
-#
-# This block reappears only when the rebase targets a vLLM commit that is NOT a
-# released tag. During the 0.27 rebase, targeting v0.27.0rc1 on top of the
-# v0.26.0 base image required reinstalling vLLM from wheels.vllm.ai and then
-# hand-repairing the fallout: pinning torch ahead of flashinfer (installing it
-# after bound flashinfer's TVM FFI runtime to the old torch and aborted engine
-# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
-# under numba's ceiling. All of it was compensation for a base image and a vLLM
-# that disagreed. At a real release boundary they agree, so none of it is
-# needed. If you pin an untagged commit again, restore the reinstall block —
-# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the
-# Dockerfile silently ship the wrong vLLM.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.27.0
+ARG VLLM_BASE_TAG=v0.26.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents

docs/user_guide/diffusion/lora.md

@@ -1,18 +1,14 @@
 # LoRA (Low-Rank Adaptation) Guide
 
-LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni supports two LoRA backends: **PEFT** for PEFT-style adapters, and **Distill** usually for few-steps inference. **PEFT** backend allows you to customize model behavior without modifying the base model weights. **Distill** backend fuses LoRA weights into the base model at initialization.
+LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights.
 
 ## Overview
 
-LoRA adapters are lightweight, model-specific fine-tuning weights that can be applied to diffusion models in two ways:
-
-- **PEFT backend** (`--lora-backend peft`, default): Loads a PEFT-format adapter folder via `DiffusionLoRAManager`. Adapters are cached (LRU) and activated per request via `LoRARequest`. It uses a unified LoRA handling mechanisms similar to vLLM with LRU cache management.
-- **Distill backend** (`--lora-backend distill`): Calls `pipeline.load_lora_weights` once at initialization to fuse one or more concrete checkpoint files directly into the base weights. Typically used for distilled few-step LoRAs (e.g. Lightning, LightX2V).
+LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management.
 
 ## LoRA Adapter Format
 
-### PEFT (Parameter-Efficient Fine-Tuning) format (default)
-A typical PEFT-format LoRA adapter directory structure:
+LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
 
 ```
 lora_adapter/
@@ -29,20 +25,19 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
 
 ### Offline Inference
 
-#### PEFT backend: pre-loaded LoRA
+#### Pre-loaded LoRA
 
-Load a PEFT-format LoRA adapter at initialization. The adapter is pre-loaded into the cache and can be activated per request:
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
 
 ```python
 from vllm_omni import Omni
 from vllm_omni.lora.request import LoRARequest
 
-lora_path = "/path/to/lora_adapter"
+lora_path="/path/to/lora_adapter"
 
 omni = Omni(
     model="stabilityai/stable-diffusion-3.5-medium",
-    lora_path=lora_path,
-    lora_backend="peft",  # default, can be omitted
+    lora_path=lora_path
 )
 
 lora_request = LoRARequest(
@@ -58,98 +53,6 @@ outputs = omni.generate(
 )
 ```
 
-#### Distill backend: fuse distilled LoRA at init
-
-For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or more concrete `.safetensors` files. The weights are fused into the base model once at init; subsequent `generate()` calls do not need a `LoRARequest`.
-
-##### Supported pipelines
-
-For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
-
-| Pipeline | Supported distilled LoRA repo | Notes |
-|----------|--------------------------------|-------|
-| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
-| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
-| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
-
-Other diffusion pipelines are not currently listed as supporting distilled LoRA. Use the PEFT backend for request-time adapters, or bake converted weights into a local Diffusers directory before serving.
-
-Single-file example (Qwen-Image-Lightning):
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Qwen/Qwen-Image-2512",
-    lora_path="/path/to/Qwen-Image-2512-Lightning.safetensors",
-    lora_backend="distill",
-)
-
-outputs = omni.generate(prompt="A piece of cheesecake")
-```
-
-Multi-file example (Wan2.2 MoE, high + low noise):
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    lora_path=[
-        "/path/to/wan2.2_high_noise_lora.safetensors",   # -> transformer
-        "/path/to/wan2.2_low_noise_lora.safetensors",    # -> transformer_2
-    ],
-    lora_backend="distill",
-)
-```
-
-The CLI examples under `examples/offline_inference/` accept the same flags, e.g.:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-  --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --lora-backend distill \
-  --lora-path /path/to/high.safetensors /path/to/low.safetensors \
-  --prompt "A cat playing with yarn"
-```
-
-##### Online serving
-
-Distilled LoRAs can also be fused when an online diffusion server starts. The
-adapter remains active for the lifetime of that server, so requests do not pass
-a per-request `lora` field.
-
-```bash
-vllm serve Qwen/Qwen-Image-2512 \
-  --omni \
-  --port 8091 \
-  --lora-backend distill \
-  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
-```
-
-Then send a normal generation request with the sampling settings expected by
-the distilled checkpoint:
-
-```bash
-curl -s http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [
-      {"role": "user", "content": "A Chinese female college student, around 20 years old, with a very short haircut that conveys a gentle, artistic vibe. Her hair naturally falls to partially cover her cheeks, projecting a tomboyish yet charming demeanor. She has cool-toned fair skin and delicate features, with a slightly shy yet subtly confident expression—her mouth crooked in a playful, youthful smirk. She wears an off-shoulder top, revealing one shoulder, with a well-proportioned figure. The image is framed as a close-up selfie: she dominates the foreground, while the background clearly shows her dormitory—a neatly made bed with white linens on the top bunk, a tidy study desk with organized stationery, and wooden cabinets and drawers. The photo is captured on a smartphone under soft, even ambient lighting, with natural tones, high clarity, and a bright, lively atmosphere full of youthful, everyday energy."}
-    ],
-    "extra_body": {
-      "height": 1024,
-      "width": 1024,
-      "num_inference_steps": 4,
-      "true_cfg_scale": 1.0,
-      "seed": 42
-    }
-  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > qwen_image_distill_lora_output.png
-```
-
-For Wan2.2 MoE serving, pass the high-noise checkpoint first and the low-noise
-checkpoint second to `--lora-path`.
-
 !!! note "Server-side Path Requirement"
     The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server.
 
@@ -238,7 +141,6 @@ Notes:
 
 - This route avoids runtime LoRA loading changes in vLLM-Omni when you choose to bake converted weights into a local Diffusers directory.
 - Output quality and speed depend on the replacement checkpoints and sampling params you choose.
-- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path <high>.safetensors <low>.safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
 
 
 ## See Also

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -88,9 +88,8 @@ python text_to_image.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
-| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 
 **NextStep-1.1 specific arguments:**
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -125,9 +125,6 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
-- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
-- `--lora-scale`: scale factor for LoRA weights.
-- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 
 ### Wan2.2-specific
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -60,10 +60,10 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.lora.request import LoRARequest
-from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import (
-    build_image_to_video_prompt,
+    build_image_to_video_prompt as build_model_image_to_video_prompt,
+)
+from vllm_omni.model_extras import (
     get_extra_body_params,
     get_model_class_name,
 )
@@ -85,6 +85,22 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
+def build_image_to_video_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    media_inputs: dict[str, Any],
+) -> dict[str, Any]:
+    """Build the canonical request envelope for the shared I2V example."""
+    result: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["video"],
+        "multi_modal_data": media_inputs,
+    }
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description=(
@@ -296,30 +312,6 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
-    parser.add_argument(
-        "--lora-path",
-        type=str,
-        nargs="+",
-        default=None,
-        help="Path to a LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. "
-        "For Wan2.2 MoE models, pass the high-noise checkpoint first and the low-noise checkpoint second.",
-    )
-    parser.add_argument(
-        "--lora-scale",
-        type=float,
-        default=1.0,
-        help="Scale factor for PEFT LoRA weights (default: 1.0).",
-    )
-    parser.add_argument(
-        "--lora-backend",
-        type=str,
-        default="peft",
-        choices=["peft", "distill"],
-        help=(
-            "LoRA backend. 'peft' loads a PEFT-format adapter for request-time activation; "
-            "'distill' fuses one or more concrete LoRA checkpoint files at initialization."
-        ),
-    )
     parser.add_argument(
         "--extra-body",
         type=functools.partial(parse_json_object, flag_name="--extra-body"),
@@ -513,12 +505,6 @@ def main():
         omni_kwargs["flow_shift"] = flow_shift
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
-    if args.lora_path is not None:
-        lora_path = args.lora_path
-        if len(lora_path) == 1:
-            lora_path = lora_path[0]
-        omni_kwargs["lora_path"] = lora_path
-        omni_kwargs["lora_backend"] = args.lora_backend
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
     if args.extra_body and "guardrails" in args.extra_body:
@@ -549,36 +535,22 @@ def main():
     print(f"  Video size: {width}x{height}")
     print(f"{'=' * 60}\n")
 
-    lora_request = None
-    if args.lora_path and args.lora_backend == "peft":
-        if len(args.lora_path) != 1:
-            raise ValueError("Only one LoRA path is expected for PEFT backend.")
-
-        lora_path = args.lora_path[0]
-        lora_request = LoRARequest(
-            lora_name=Path(lora_path).stem,
-            lora_int_id=stable_lora_int_id(lora_path),
-            lora_path=lora_path,
-        )
-
     negative_prompt = args.negative_prompt
     if negative_prompt is None and not is_ltx2:
         # Preserve the historical empty-prompt behavior for non-LTX examples.
         negative_prompt = ""
     prompt_dict = build_image_to_video_prompt(
-        model_class_name=model_class_name,
         prompt=args.prompt,
         negative_prompt=negative_prompt,
         media_inputs=media_inputs,
+    )
+    prompt_dict = build_model_image_to_video_prompt(
+        model_class_name=model_class_name,
+        prompt=prompt_dict,
         height=height,
         width=width,
         num_frames=num_frames,
     )
-    sampling_extra_args = {"sample_solver": args.sample_solver}
-    if lora_request is not None:
-        sampling_extra_args["lora_request"] = lora_request
-        sampling_extra_args["lora_scale"] = args.lora_scale
-
     sampling_params = OmniDiffusionSamplingParams(
         height=height,
         width=width,
@@ -589,7 +561,7 @@ def main():
         num_inference_steps=num_inference_steps,
         num_frames=num_frames,
         frame_rate=frame_rate,
-        extra_args=sampling_extra_args,
+        extra_args={"sample_solver": args.sample_solver},
     )
     if flow_shift is not None:
         sampling_params.extra_args["flow_shift"] = flow_shift

examples/offline_inference/text_to_image/README.md

@@ -57,8 +57,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output lingbot_t2i.png
 ```
 
-The LingBot prompt builder selects the `image` output modality and requests one
-frame. Passing an input image is not supported in LingBot T2I mode.
+The shared T2I entry point selects the `image` output modality. The LingBot
+pipeline validates that T2I requests produce one frame and do not include an
+input image.
 
 ## Quick Start
 
@@ -109,9 +110,8 @@ python text_to_image.py \
 | `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
 | `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
-| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
 | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 | `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |

examples/offline_inference/text_to_image/text_to_image.py

@@ -21,7 +21,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import (
-    build_text_to_image_prompt,
+    build_text_to_image_prompt as build_model_text_to_image_prompt,
+)
+from vllm_omni.model_extras import (
     get_extra_body_params,
     get_model_class_name,
     should_init_extra_args_for_non_diffusion_stages,
@@ -56,6 +58,17 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
+def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
+    """Build the canonical request envelope for the shared T2I example."""
+    result: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["image"],
+    }
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result
+
+
 def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     """Convert NumPy diffusion outputs to PIL images before saving."""
     normalized = []
@@ -260,7 +273,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--lora-path",
         type=str,
-        nargs="+",
         default=None,
         help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.",
     )
@@ -270,15 +282,6 @@ def parse_args() -> argparse.Namespace:
         default=1.0,
         help="Scale factor for LoRA weights (default: 1.0).",
     )
-    parser.add_argument(
-        "--lora-backend",
-        type=str,
-        default="peft",
-        choices=["peft", "distill"],
-        help="LoRA backend for loading LoRA adapters. Default: peft"
-        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
-        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
-    )
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
@@ -414,12 +417,8 @@ def main():
     # Prepare LoRA kwargs for Omni initialization
     lora_args: dict[str, Any] = {}
     if args.lora_path:
-        lora_path = args.lora_path
-        if len(lora_path) == 1:
-            lora_path = lora_path[0]
-        lora_args["lora_path"] = lora_path
-        lora_args["lora_backend"] = args.lora_backend
-        print(f"Using LoRA from: {lora_path} with backend: {args.lora_backend}")
+        lora_args["lora_path"] = args.lora_path
+        print(f"Using LoRA from: {args.lora_path}")
 
     # Build quantization kwargs: use quantization_config dict when
     # ignored_layers is specified so the list flows through OmniDiffusionConfig
@@ -508,24 +507,23 @@ def main():
 
     # Build LoRA request when --lora-path is set
     lora_request = None
-    if args.lora_path and args.lora_backend == "peft":
-        if len(args.lora_path) != 1:
-            raise ValueError("Only one LoRA path is expected for PEFT backend.")
-
-        lora_path = args.lora_path[0]
-        lora_request_id = stable_lora_int_id(lora_path)
+    if args.lora_path:
+        lora_request_id = stable_lora_int_id(args.lora_path)
         lora_request = LoRARequest(
-            lora_name=Path(lora_path).stem,
+            lora_name=Path(args.lora_path).stem,
             lora_int_id=lora_request_id,
-            lora_path=lora_path,
+            lora_path=args.lora_path,
         )
 
     generation_start = time.perf_counter()
 
     prompt_dict = build_text_to_image_prompt(
-        model_class_name=model_class_name,
         prompt=args.prompt,
         negative_prompt=args.negative_prompt,
+    )
+    prompt_dict = build_model_text_to_image_prompt(
+        model_class_name=model_class_name,
+        prompt=prompt_dict,
         height=args.height,
         width=args.width,
     )

examples/offline_inference/text_to_video/text_to_video.md

@@ -218,9 +218,6 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
-- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
-- `--lora-scale`: scale factor for LoRA weights.
-- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 
 ### Wan2.2-specific

examples/offline_inference/text_to_video/text_to_video.py

@@ -14,8 +14,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.lora.request import LoRARequest
-from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -336,29 +334,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
-    parser.add_argument(
-        "--lora-path",
-        type=str,
-        nargs="+",
-        default=None,
-        help="Path to LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. Loaded at initialization and used for generation."
-        "Note: for Wan2.2 MoE models, two checkpoints positionally mapped to high-noise and low-noise modules.",
-    )
-    parser.add_argument(
-        "--lora-scale",
-        type=float,
-        default=1.0,
-        help="Scale factor for LoRA weights (default: 1.0).",
-    )
-    parser.add_argument(
-        "--lora-backend",
-        type=str,
-        default="peft",
-        choices=["peft", "distill"],
-        help="LoRA backend for loading LoRA adapters. Default: peft"
-        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
-        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
-    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -465,12 +440,6 @@ def main():
         omni_kwargs["cache_backend"] = args.cache_backend
         omni_kwargs["cache_config"] = cache_config
         omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
-    if args.lora_path is not None:
-        lora_path = args.lora_path
-        if len(lora_path) == 1:
-            lora_path = lora_path[0]
-        omni_kwargs["lora_path"] = lora_path
-        omni_kwargs["lora_backend"] = args.lora_backend
 
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
@@ -500,19 +469,6 @@ def main():
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
-    lora_request = None
-    if args.lora_path and args.lora_backend == "peft":
-        if len(args.lora_path) != 1:
-            raise ValueError("Only one LoRA path is expected for PEFT backend.")
-
-        lora_path = args.lora_path[0]
-        lora_request_id = stable_lora_int_id(lora_path)
-        lora_request = LoRARequest(
-            lora_name=Path(lora_path).stem,
-            lora_int_id=lora_request_id,
-            lora_path=lora_path,
-        )
-
     prompt_dict = {"prompt": args.prompt}
     if args.negative_prompt is not None:
         prompt_dict["negative_prompt"] = args.negative_prompt
@@ -520,11 +476,6 @@ def main():
         # Preserve the historical empty-prompt behavior for non-LTX examples.
         prompt_dict["negative_prompt"] = ""
 
-    extra_args = {}
-    if lora_request:
-        extra_args["lora_request"] = lora_request
-        extra_args["lora_scale"] = args.lora_scale
-
     sampling_kwargs = dict(
         height=args.height,
         width=args.width,
@@ -532,7 +483,6 @@ def main():
         guidance_scale=args.guidance_scale,
         num_inference_steps=args.num_inference_steps,
         num_frames=args.num_frames,
-        extra_args=extra_args,
     )
     if args.guidance_scale_high is not None:
         sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high

examples/online_serving/image_to_video/README.md

@@ -1,6 +1,9 @@
 # Image-To-Video
 
-This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
+This example demonstrates online image-to-video generation with vLLM-Omni.
+The startup script defaults to Wan2.2, but `MODEL` can select any supported
+image-to-video model. The existing curl helper remains a concrete Wan2.2
+request through the standard Videos API.
 
 ## Start Server
 
@@ -20,9 +23,7 @@ bash run_server.sh
 
 The script allows overriding:
 - `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
-- `PORT` (default: `8091`)
-- `BOUNDARY_RATIO` (default: `0.875`)
-- `FLOW_SHIFT` (default: `12.0`)
+- `PORT` (default: `8099`)
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 

examples/online_serving/image_to_video/run_server.sh

@@ -1,12 +1,12 @@
 #!/bin/bash
-# Wan2.2 image-to-video server start script
+# Shared image-to-video server start script
 
 MODEL="${MODEL:-Wan-AI/Wan2.2-I2V-A14B-Diffusers}"
 PORT="${PORT:-8099}"
 CACHE_BACKEND="${CACHE_BACKEND:-none}"
 ENABLE_CACHE_DIT_SUMMARY="${ENABLE_CACHE_DIT_SUMMARY:-0}"
 
-echo "Starting Wan2.2 I2V server..."
+echo "Starting image-to-video server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
 echo "Cache backend: $CACHE_BACKEND"

examples/online_serving/text_to_image/README.md

@@ -1,6 +1,9 @@
 # Text-To-Image
 
-This example demonstrates how to deploy Qwen-Image model for online image generation service using vLLM-Omni.
+This example demonstrates online text-to-image generation with vLLM-Omni. The
+startup script and Python clients default to Qwen-Image, but `MODEL` can select
+any supported text-to-image model. The existing curl helper remains a concrete
+Qwen-Image request through the standard Images API.
 
 ## Start Server
 
@@ -115,23 +118,7 @@ python gradio_demo.py
 
 ## LoRA
 
-This example supports both startup-time distilled LoRAs and request-time
-PEFT-compatible adapters. The LoRA path must be readable on the **server**
-machine (usually a local path or a mounted directory).
-
-### Start Server with a Distilled LoRA
-
-```bash
-vllm serve Qwen/Qwen-Image-2512 \
-  --omni \
-  --port 8091 \
-  --lora-backend distill \
-  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
-```
-
-The distilled LoRA is fused once during server initialization and applies to
-every request. Use the checkpoint's few-step sampling settings, for example
-`num_inference_steps=4` and `true_cfg_scale=1.0`.
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
 
 ### Using Python Client with LoRA
 
@@ -242,8 +229,8 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `height`                 | int   | None    | Image height in pixels         |
 | `width`                  | int   | None    | Image width in pixels          |
 | `size`                   | str   | None    | Image size (e.g., "1024x1024") |
-| `num_inference_steps`    | int   | 50      | Number of denoising steps      |
-| `true_cfg_scale`         | float | 4.0     | Qwen-Image CFG scale           |
+| `num_inference_steps`    | int   | model   | Number of denoising steps      |
+| `true_cfg_scale`         | float | model   | True CFG scale when supported  |
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |

examples/online_serving/text_to_image/gradio_demo.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Qwen-Image Gradio Demo for online serving.
+Shared text-to-image Gradio demo for online serving.
 
 Usage:
     python gradio_demo.py [--server http://localhost:8091] [--port 7860]
@@ -77,9 +77,9 @@ def generate_image(
 def create_demo(server_url: str):
     """Create Gradio demo interface."""
 
-    with gr.Blocks(title="Qwen-Image Demo") as demo:
-        gr.Markdown("# Qwen-Image Online Generation")
-        gr.Markdown("Generate images using Qwen-Image model")
+    with gr.Blocks(title="Text-to-Image Demo") as demo:
+        gr.Markdown("# Text-to-Image Online Generation")
+        gr.Markdown("Generate images using the model configured on the server")
 
         with gr.Row():
             with gr.Column(scale=1):
@@ -115,7 +115,6 @@ def create_demo(server_url: str):
                         label="Inference Steps",
                         minimum=10,
                         maximum=100,
-                        # Default steps aligned with run_curl_text_to_image.sh to 100
                         value=100,
                         step=5,
                     )
@@ -173,7 +172,7 @@ def create_demo(server_url: str):
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Qwen-Image Gradio Demo")
+    parser = argparse.ArgumentParser(description="Shared text-to-image Gradio demo")
     parser.add_argument("--server", default="http://localhost:8091", help="Server URL")
     parser.add_argument("--port", type=int, default=7860, help="Gradio port")
     parser.add_argument("--share", action="store_true", help="Create public link")

examples/online_serving/text_to_image/openai_chat_client.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Qwen-Image OpenAI-compatible image generation client.
+Shared OpenAI-compatible text-to-image generation client.
 
 Usage:
     python openai_chat_client.py --prompt "A beautiful landscape" --output output.png
@@ -39,7 +39,7 @@ def generate_image(
         height: Image height in pixels
         width: Image width in pixels
         steps: Number of diffusion steps
-        true_cfg_scale: Qwen-Image CFG scale
+        true_cfg_scale: True CFG scale for models that support it
         seed: Random seed
         negative_prompt: Negative prompt
         num_outputs_per_prompt: Number of images to generate
@@ -70,7 +70,7 @@ def generate_image(
         payload["num_inference_steps"] = steps
     if true_cfg_scale is not None:
         payload["true_cfg_scale"] = true_cfg_scale
-    if negative_prompt:
+    if negative_prompt is not None:
         payload["negative_prompt"] = negative_prompt
     if seed is not None:
         payload["seed"] = seed
@@ -115,15 +115,15 @@ def generate_image(
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Qwen-Image chat client")
+    parser = argparse.ArgumentParser(description="Shared text-to-image client")
     parser.add_argument("--prompt", "-p", default="a cup of coffee on the table", help="Text prompt")
-    parser.add_argument("--output", "-o", default="qwen_image_output.png", help="Output file")
+    parser.add_argument("--output", "-o", default="text_to_image_output.png", help="Output file")
     parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
-    parser.add_argument("--height", type=int, default=1024, help="Image height")
-    parser.add_argument("--width", type=int, default=1024, help="Image width")
-    parser.add_argument("--steps", type=int, default=50, help="Inference steps")
-    parser.add_argument("--cfg-scale", type=float, default=4.0, help="True CFG scale")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--height", type=int, help="Image height; defaults to the model setting")
+    parser.add_argument("--width", type=int, help="Image width; defaults to the model setting")
+    parser.add_argument("--steps", type=int, help="Inference steps; defaults to the model setting")
+    parser.add_argument("--cfg-scale", type=float, help="True CFG scale when supported")
+    parser.add_argument("--seed", type=int, help="Random seed")
     parser.add_argument("--negative", help="Negative prompt")
 
     parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")

examples/online_serving/text_to_image/run_server.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
-# Qwen-Image online serving startup script
+# Shared text-to-image online serving startup script
 
 MODEL="${MODEL:-Qwen/Qwen-Image}"
 PORT="${PORT:-8091}"
 
-echo "Starting Qwen-Image server..."
+echo "Starting text-to-image server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
 

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -92,9 +92,6 @@ class _Request:
         self.num_computed_tokens = 0
         self.num_in_flight_tokens = 0
         self.num_output_placeholders = 0
-        # vLLM 0.27 (a0c092ee72): Request gained num_stale_output_tokens to
-        # track in-flight outputs discarded at preemption/streaming-stop.
-        self.num_stale_output_tokens = 0
         self.has_encoder_inputs = False
         self.pooling_params = None
         self.resumable = True

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -1,223 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Investigation tests for the num_stale_output_tokens drain (59b9e719).
-
-The seed-TTS WER regression appeared with the vLLM 0.27 rebase (nightly
-2946 -> 2949: mean WER 0.2695 -> 0.3936, median 0.0 in both), whose prime
-suspect is the reimplemented async-token discard: vLLM 0.26 drained
-``async_tokens_to_discard`` one arriving frame at a time and suppressed
-only the token append; the 0.27 adaptation seeds
-``num_stale_output_tokens`` with the discarded *placeholder count* and
-drains it by each arriving frame's ``num_tokens_scheduled``, dropping the
-whole per-request update. Whenever seed and arrivals disagree, the counter
-outlives the old segment and the drain consumes *valid* frames of the new
-segment — a hole in the talker's codec stream (garbled audio, text/audio
-mismatch) — or underflows its own assert on the new segment's prefill
-frame (engine-core death mid-step).
-
-The fix seeds the counter from ``num_in_flight_tokens`` — the same
-scheduled-token units the drain subtracts — so every pre-discard frame
-drains exactly its own contribution and new-segment frames can never be
-swallowed. These tests pin that contract one frame at a time; the
-delivered / swallowed observable is whether ``_update_request_with_output``
-runs for the frame.
-"""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-from unittest.mock import MagicMock
-
-import pytest
-
-# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
-# before Request / StreamingUpdate are bound in this module.
-# isort: off
-import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
-from vllm.sampling_params import SamplingParams
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.outputs import ModelRunnerOutput
-from vllm.v1.request import Request, RequestStatus, StreamingUpdate
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
-# isort: on
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_session() -> Request:
-    session = Request(
-        request_id="req-stale-drain-test",
-        prompt_token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=32),
-        pooling_params=None,
-        arrival_time=100.0,
-        block_hasher=None,
-    )
-    session.status = RequestStatus.RUNNING
-    if not hasattr(session, "num_stale_output_tokens"):
-        # vLLM 0.27 defines this Request field; 0.26 (local dev) does not.
-        session.num_stale_output_tokens = 0
-    if not hasattr(session, "num_in_flight_tokens"):
-        session.num_in_flight_tokens = 0
-    return session
-
-
-def _replace_streaming_session(session: Request) -> None:
-    """Run the real stage-0 streaming replacement (the discard/seed site)."""
-    sched = OmniARScheduler.__new__(OmniARScheduler)
-    sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=0))
-    sched.num_waiting_for_streaming_input = 0
-    sched.log_stats = False
-    sched.chunk_transfer_adapter = None
-    sched.skipped_waiting = set()
-    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    update = StreamingUpdate(
-        mm_features=None,
-        prompt_token_ids=[10, 20],
-        max_tokens=32,
-        arrival_time=200.0,
-        sampling_params=SamplingParams(max_tokens=16),
-    )
-    sched._update_request_as_session(session, update)
-    session.status = RequestStatus.RUNNING
-
-
-def _make_drain_sched(session: Request) -> MagicMock:
-    sched = MagicMock()
-    sched.requests = {session.request_id: session}
-    sched.perf_metrics = None
-    sched.structured_output_manager.should_advance.return_value = False
-    sched._update_request_with_output.return_value = ([42], False)
-    sched._process_kv_transfer_trigger.return_value = False
-    sched.chunk_transfer_adapter = MagicMock()
-    sched.running = [session]
-    sched.waiting_for_transfer_free = set()
-    sched.transfer_triggered_requests = set()
-    sched.active_kv_transfers = set()
-    sched.pending_stop_after_extraction = set()
-    sched.connector = None
-    sched.kv_cache_manager.take_events.return_value = None
-    sched.finished_req_ids_dict = {}
-    sched.make_stats.return_value = None
-    return sched
-
-
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
-    """Feed one model-runner frame through update_from_output.
-
-    Returns True when the frame's tokens were delivered (the append ran),
-    False when the frame was swallowed by the stale drain.
-    """
-    scheduler_output = MagicMock(spec=SchedulerOutput)
-    scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
-    scheduler_output.num_invalid_spec_tokens = 0
-
-    model_runner_output = MagicMock(spec=ModelRunnerOutput)
-    model_runner_output.sampled_token_ids = [[token]]
-    model_runner_output.logprobs = None
-    model_runner_output.prompt_logprobs_dict = {}
-    model_runner_output.pooler_output = None
-    model_runner_output.num_nans_in_logits = None
-    model_runner_output.kv_connector_output = None
-    model_runner_output.cudagraph_stats = None
-    model_runner_output.req_id_to_index = {session.request_id: 0}
-    model_runner_output.routed_experts = None
-    model_runner_output.inter_stage_outputs = None
-
-    sched._update_request_with_output.reset_mock()
-    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
-    return sched._update_request_with_output.called
-
-
-def test_exact_drain_delivers_new_segment_frame() -> None:
-    """One in-flight decode frame at replacement: its late output drains the
-    counter exactly and the new segment's first frame is delivered."""
-    session = _make_session()
-    session.append_output_token_ids([7, 8, 9])
-    session.num_computed_tokens = 6
-    session.num_output_placeholders = 1
-    session.num_in_flight_tokens = 1
-
-    _replace_streaming_session(session)
-    assert session.num_stale_output_tokens == 1
-
-    sched = _make_drain_sched(session)
-    assert _run_step(sched, session, num_scheduled=1, token=42) is False  # late frame dropped
-    assert session.num_stale_output_tokens == 0
-    assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
-
-
-def test_two_in_flight_frames_drain_without_swallowing() -> None:
-    """Two in-flight decode frames: both late frames drop, the new segment's
-    first valid frame is delivered. A placeholder-based seed swallowed it
-    whenever fewer frames arrived than placeholders counted (the driver of
-    the seed-TTS WER regression, nightly 2946 -> 2949: 0.2695 -> 0.3936)."""
-    session = _make_session()
-    session.append_output_token_ids([7, 8, 9])
-    session.num_computed_tokens = 7
-    session.num_output_placeholders = 2
-    session.num_in_flight_tokens = 2
-
-    _replace_streaming_session(session)
-    assert session.num_stale_output_tokens == 2
-
-    sched = _make_drain_sched(session)
-    assert _run_step(sched, session, num_scheduled=1, token=42) is False  # late frame 1
-    assert _run_step(sched, session, num_scheduled=1, token=44) is False  # late frame 2
-    assert session.num_stale_output_tokens == 0
-
-    delivered = _run_step(sched, session, num_scheduled=1, token=43)  # new segment, frame 1
-    assert delivered is True, (
-        "valid new-segment frame was swallowed by leftover stale count "
-        f"(num_stale_output_tokens left at {session.num_stale_output_tokens})"
-    )
-
-
-def test_spec_placeholder_asymmetry_does_not_swallow_valid_frame() -> None:
-    """Placeholder counts can exceed scheduled counts (e.g. a draft
-    invalidated before execution): 1 in-flight MTP frame scheduled 2 tokens
-    but contributed 3 placeholders. Seeding from num_in_flight_tokens keeps
-    seed == drain, so the new segment's valid frame is delivered."""
-    session = _make_session()
-    session.append_output_token_ids([7])
-    session.num_computed_tokens = 7
-    session.num_output_placeholders = 3  # 1 sampled + 2 spec placeholders
-    session.num_in_flight_tokens = 2  # the frame actually scheduled 1 + 1 tokens
-
-    _replace_streaming_session(session)
-    assert session.num_stale_output_tokens == 2
-
-    sched = _make_drain_sched(session)
-    assert _run_step(sched, session, num_scheduled=2, token=42) is False  # late MTP frame
-    assert session.num_stale_output_tokens == 0
-    delivered = _run_step(sched, session, num_scheduled=1, token=43)  # new segment, frame 1
-    assert delivered is True, (
-        "valid new-segment frame swallowed after spec placeholder asymmetry "
-        f"(num_stale_output_tokens left at {session.num_stale_output_tokens})"
-    )
-
-
-def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
-    """An in-flight prefill chunk carries no placeholders but 5 scheduled
-    tokens. The placeholder-based seed missed it entirely (its late output
-    was appended to the rolled-back session) and any leftover counter
-    underflowed 'assert num_stale_output_tokens >= 0' on the chunk's arrival,
-    killing the engine-core step. Seeding from num_in_flight_tokens drops the
-    late chunk and drains to exactly zero."""
-    session = _make_session()
-    session.append_output_token_ids([7, 8, 9])
-    session.num_computed_tokens = 7
-    session.num_output_placeholders = 0  # prefill chunk: no placeholders
-    session.num_in_flight_tokens = 5  # the chunk's scheduled tokens
-
-    _replace_streaming_session(session)
-    assert session.num_stale_output_tokens == 5
-
-    sched = _make_drain_sched(session)
-    assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
-    assert session.num_stale_output_tokens == 0
-    assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -64,7 +64,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         local_rank,
         backend,
         group_name,
-        use_all2all=False,
     ):
         del backend
         group = _FakeGroup(
@@ -74,7 +73,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
             device_communicator=object(),
             reduce_scatter=lambda tensor, **kwargs: tensor,
         )
-        group.use_all2all = use_all2all
         created_groups.append(group)
         return group
 
@@ -123,9 +121,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
     assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
     assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
-    # vLLM 0.27's MoE oracle asserts the EP communicator has an all2all
-    # manager, which vLLM only builds for groups created with use_all2all.
-    assert parallel_state.vllm_parallel_state._EP.use_all2all is True
     assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
     assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
     assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
@@ -318,35 +313,3 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
     assert parallel_state.vllm_parallel_state._PCP is None
     assert parallel_state.vllm_parallel_state._EP is None
     assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
-
-
-@pytest.mark.cpu
-@pytest.mark.core_model
-def test_init_vllm_group_forwards_use_all2all_only_when_supported(monkeypatch):
-    """use_all2all reaches vLLM only on versions whose API accepts it.
-
-    vLLM 0.27 builds the EP all2all manager only for groups created with
-    use_all2all=True (its MoE oracle then asserts the manager exists —
-    build 2953, HunyuanImage3-DIT accuracy). Older vLLM has no such kwarg,
-    so forwarding it unconditionally would TypeError at init.
-    """
-    received = {}
-
-    def modern(group_ranks, local_rank, backend, group_name=None, use_device_communicator=True, use_all2all=False):
-        received["modern"] = use_all2all
-        return object()
-
-    monkeypatch.setattr(parallel_state.vllm_parallel_state, "init_model_parallel_group", modern)
-    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "ep", use_all2all=True)
-    assert received["modern"] is True
-
-    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "tp")
-    assert received["modern"] is False
-
-    def legacy(group_ranks, local_rank, backend, group_name=None, use_device_communicator=True):
-        received["legacy"] = True
-        return object()
-
-    monkeypatch.setattr(parallel_state.vllm_parallel_state, "init_model_parallel_group", legacy)
-    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "ep", use_all2all=True)
-    assert received["legacy"] is True

tests/diffusion/layers/test_fused_moe.py

@@ -110,7 +110,7 @@ class TestFusedMoEFactory:
         monkeypatch.setitem(
             sys.modules,
             "vllm.model_executor.layers.fused_moe",
-            SimpleNamespace(FusedMoEFactory=mock_vllm_fused_moe),
+            SimpleNamespace(FusedMoE=mock_vllm_fused_moe),
         )
 
         result = fused_moe.FusedMoE(prefix="test", a=1, reduce_results=True)
@@ -139,7 +139,7 @@ class TestFusedMoEFactory:
         monkeypatch.setitem(
             sys.modules,
             "vllm.model_executor.layers.fused_moe",
-            SimpleNamespace(FusedMoEFactory=mocker.MagicMock(return_value=mock_runner)),
+            SimpleNamespace(FusedMoE=mocker.MagicMock(return_value=mock_runner)),
         )
         mocker.patch.object(fused_moe, "current_omni_platform", mocker.MagicMock())
         mock_set_num_tokens = mocker.patch.object(fused_moe, "_set_forward_context_num_tokens")
@@ -169,7 +169,7 @@ class TestFusedMoEFactory:
         monkeypatch.setitem(
             sys.modules,
             "vllm.model_executor.layers.fused_moe",
-            SimpleNamespace(FusedMoEFactory=mocker.MagicMock(return_value=mock_runner)),
+            SimpleNamespace(FusedMoE=mocker.MagicMock(return_value=mock_runner)),
         )
         mocker.patch.object(fused_moe, "current_omni_platform", mocker.MagicMock())
         mock_set_num_tokens = mocker.patch.object(fused_moe, "_set_forward_context_num_tokens")

tests/diffusion/lora/test_loader.py

@@ -1,438 +0,0 @@
-from collections import defaultdict
-
-import pytest
-import torch
-import torch.nn as nn
-from diffusers.loaders.lora_conversion_utils import (
-    _convert_non_diffusers_qwen_lora_to_diffusers,
-)
-from pytest_mock import MockerFixture
-from torch.testing import assert_close
-
-from vllm_omni.diffusion.lora.loader import (
-    LoraLoaderMixin,
-    QwenImageLoraLoaderMixin,
-    WanLoraLoaderMixin,
-    _prepare_lora_delta,
-    _remap_state_dict_keys,
-)
-
-HEAD_DIM = 24
-RANK = 2
-NUM_LAYERS = 10
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class DummyTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        d_model: int,
-        bias: bool = False,
-        stacked_params: bool = False,
-    ):
-        super().__init__()
-        self.attn = nn.Module()
-        if stacked_params:
-            self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias)
-        else:
-            self.attn.to_q = nn.Linear(d_model, d_model, bias=bias)
-            self.attn.to_k = nn.Linear(d_model, d_model, bias=bias)
-            self.attn.to_v = nn.Linear(d_model, d_model, bias=bias)
-        self.attn.to_out = nn.Linear(d_model, d_model, bias=bias)
-
-
-class DummyTransformer(nn.Module):
-    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
-        super().__init__()
-
-        self.blocks = nn.ModuleList(
-            [DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)]
-        )
-        if stacked_params:
-            self.stacked_params_mapping = [
-                (".attn.to_qkv", ".attn.to_q", "q"),
-                (".attn.to_qkv", ".attn.to_k", "k"),
-                (".attn.to_qkv", ".attn.to_v", "v"),
-            ]
-
-
-class DummyPipeline(nn.Module, LoraLoaderMixin):
-    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
-        super().__init__()
-        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params)
-
-
-# ======================================================
-# helper functions for test
-# ======================================================
-def _filter_parameters(named_params, mode: str = "all"):
-    if mode == "all":
-        return named_params
-
-    filtered = []
-    for name, param in named_params:
-        if name.endswith(mode):
-            filtered.append((name, param))
-
-    return filtered
-
-
-def assert_params_equal(module, original_params, mode: str = "all"):
-    for name, param in _filter_parameters(module.named_parameters(), mode):
-        assert_close(param, original_params[name])
-
-
-def assert_params_not_equal(module, original_params, mode: str = "all"):
-    for name, param in _filter_parameters(module.named_parameters(), mode):
-        assert not torch.allclose(param, original_params[name])
-
-
-def make_dummy_lora_state_dict(
-    base_key,
-    weight_a: torch.Tensor | None = None,
-    weight_b: torch.Tensor | None = None,
-    weight_bias: torch.Tensor | None = None,
-    bias: bool = False,
-    lora_a_suffix: str = "lora_A.weight",
-    lora_b_suffix: str = "lora_B.weight",
-    lora_bias_suffix: str = "bias",
-):
-    """
-    Create a dummy LoRA state dict.
-
-    Returns:
-        tuple: (state_dict, expected_delta)
-    """
-    key_a = f"{base_key}.{lora_a_suffix}"
-    key_b = f"{base_key}.{lora_b_suffix}"
-    key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
-
-    if weight_a is None:
-        weight_a = torch.randn(RANK, HEAD_DIM)
-    if weight_b is None:
-        weight_b = torch.randn(HEAD_DIM, RANK)
-
-    state_dict = {}
-    expected_delta = None
-
-    if bias:
-        state_dict[key_a] = weight_a
-        state_dict[key_b] = weight_b
-        weight_bias = torch.randn(HEAD_DIM)
-        state_dict[key_bias] = weight_bias
-        expected_delta = weight_bias
-    else:
-        state_dict[key_a] = weight_a
-        state_dict[key_b] = weight_b
-        expected_delta = weight_b @ weight_a
-
-    return state_dict, expected_delta
-
-
-def make_lora_state_dict_for_module(
-    module,
-    lora_rank=RANK,
-    bias=False,
-    lora_a_suffix: str = "lora_A.weight",
-    lora_b_suffix: str = "lora_B.weight",
-    lora_bias_suffix: str = "bias",
-    prefix: str = "transformer",
-    remap_proj_out: bool = True,
-):
-    param_to_weight_names = defaultdict(list)
-    if hasattr(module, "stacked_params_mapping"):
-        for param_name, weight_name, _ in module.stacked_params_mapping:
-            param_to_weight_names[param_name].append(weight_name)
-
-    state_dict = {}
-    for name, param in module.named_parameters(prefix=prefix):
-        if name.endswith(".weight"):
-            base_key = name[: -len(".weight")]
-            if remap_proj_out and base_key.endswith(".to_out"):
-                base_key = base_key.replace(".to_out", ".to_out.0")
-
-            d_in, d_out = param.shape
-            is_stacked_param = False
-            for param_name, weight_names in param_to_weight_names.items():
-                if param_name not in name:
-                    continue
-                is_stacked_param = True
-                target_in = d_in // len(weight_names)
-                for weight_name in weight_names:
-                    lora_key = base_key.replace(param_name, weight_name)
-                    weight_a = torch.randn(lora_rank, d_out)
-                    weight_b = torch.randn(target_in, lora_rank)
-                    state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a
-                    state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b
-
-            if is_stacked_param:
-                continue
-
-            weight_a = torch.randn(lora_rank, d_out)
-            weight_b = torch.randn(d_in, lora_rank)
-            key_a = f"{base_key}.{lora_a_suffix}"
-            key_b = f"{base_key}.{lora_b_suffix}"
-            state_dict[key_a] = weight_a
-            state_dict[key_b] = weight_b
-        elif name.endswith(".bias") and bias:
-            base_key = name[: -len(".bias")]
-            if remap_proj_out and base_key.endswith(".to_out"):
-                base_key = base_key.replace(".to_out", ".to_out.0")
-
-            d_bias = param.shape[0]
-            is_stacked_param = False
-            for param_name, weight_names in param_to_weight_names.items():
-                if param_name not in name:
-                    continue
-                is_stacked_param = True
-                target_bias = d_bias // len(weight_names)
-                for weight_name in weight_names:
-                    lora_key = base_key.replace(param_name, weight_name)
-                    weight_bias = torch.randn(target_bias)
-                    state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias
-
-            if is_stacked_param:
-                continue
-
-            key_bias = f"{base_key}.{lora_bias_suffix}"
-            state_dict[key_bias] = torch.randn_like(param)
-
-    return state_dict
-
-
-# ======================================================
-# Test loader helper functions
-# ======================================================
-def test_prepare_lora_delta_with_stacked_params():
-    base_key = "transformer.attn.to_qkv"
-
-    sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q")
-    sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k")
-    sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v")
-
-    lora_state_dict = {}
-    lora_state_dict.update(sd0)
-    lora_state_dict.update(sd1)
-    lora_state_dict.update(sd2)
-    expected_delta = torch.concat([delta0, delta1, delta2])
-
-    stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]}
-
-    actual_delta, used_keys = _prepare_lora_delta(
-        lora_state_dict,
-        base_key,
-        stacked_mapping,
-    )
-
-    # validate the delta
-    assert actual_delta is not None
-    assert_close(actual_delta, expected_delta)
-
-    # validate the used keys
-    key_q_a = "transformer.attn.to_q.lora_A.weight"
-    key_q_b = "transformer.attn.to_q.lora_B.weight"
-    key_k_a = "transformer.attn.to_k.lora_A.weight"
-    key_k_b = "transformer.attn.to_k.lora_B.weight"
-    key_v_a = "transformer.attn.to_v.lora_A.weight"
-    key_v_b = "transformer.attn.to_v.lora_B.weight"
-
-    for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]:
-        assert key in used_keys
-
-
-def test_remap_state_dict_keys_matched_two():
-    lora_state_dict, _ = make_dummy_lora_state_dict(
-        "transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias"
-    )
-
-    lora_state_dict = _remap_state_dict_keys(
-        lora_state_dict,
-        [(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")],
-    )
-
-    assert len(lora_state_dict) == 3
-    assert "transformer.attn.to_out.bias" in lora_state_dict
-
-
-# ======================================================
-# Test LoraLoaderMixin
-# ======================================================
-class TestLoraLoaderMixin:
-    @pytest.mark.parametrize(
-        "bias, stacked_params",
-        [
-            (False, False),
-            (True, False),
-            (False, True),
-            (True, True),
-        ],
-        ids=[
-            "basic",
-            "bias",
-            "stacked_params",
-            "bias_stacked_params",
-        ],
-    )
-    def test_load_lora_into_module(self, bias: bool, stacked_params: bool):
-        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params)
-        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
-
-        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias)
-        assert len(lora_state_dict) > 0
-
-        # remap the keys to vllm-omni format
-        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
-
-        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
-        # validate the parameters are updated
-        assert_params_not_equal(pipeline.transformer, original_params)
-
-        # validate all keys are used
-        missing_keys = set(lora_state_dict.keys()) - used_keys
-        assert len(missing_keys) == 0
-
-    def test_unload_module_lora(self):
-        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
-        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
-
-        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
-        # validate the state dict is valid
-        assert len(lora_state_dict) > 0
-
-        # remap the keys to vllm-omni format
-        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
-
-        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
-
-        # validate the weights are updated
-        assert_params_not_equal(pipeline.transformer, original_params)
-
-        # validate all keys are used
-        missing_keys = set(lora_state_dict.keys()) - used_keys
-        assert len(missing_keys) == 0
-
-        unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer)
-        missing_keys = set(lora_state_dict.keys()) - unload_used_keys
-        assert len(missing_keys) == 0
-        assert len(unload_used_keys - used_keys) == 0
-
-        # validate the weights are restored
-        assert_params_equal(pipeline.transformer, original_params)
-
-
-# ======================================================
-# Test QwenImageLoraLoaderMixin
-# ======================================================
-class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin):
-    def __init__(self, num_layers: int, d_model: int, bias: bool = False):
-        super().__init__()
-        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True)
-
-
-class TestQwenImageLoraLoaderMixin:
-    def test_load_lora_weights(self, mocker: MockerFixture):
-        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
-        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
-
-        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
-        assert len(lora_state_dict) > 0
-
-        mocker.patch(
-            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
-            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
-        )
-        pipeline.load_lora_weights(lora_state_dict, "adapter0")
-
-        # validate the weights are updated after lora loaded
-        assert_params_not_equal(pipeline.transformer, original_params)
-
-        # validate lora_loaded map is updated after lora loaded
-        assert "adapter0" in pipeline.lora_loaded
-
-    def test_unload_lora_weights(self, mocker: MockerFixture):
-        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
-        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
-
-        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
-        assert len(lora_state_dict) > 0
-
-        mocker.patch(
-            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
-            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
-        )
-        pipeline.load_lora_weights(lora_state_dict, "adapter0")
-        # validate the weights are updated after lora loaded
-        assert_params_not_equal(pipeline.transformer, original_params)
-
-        # validate lora_loaded map is updated after lora loaded
-        assert "adapter0" in pipeline.lora_loaded
-
-        pipeline.unload_lora_weights("adapter0")
-        # validate the weights are restored after lora unloaded
-        assert_params_equal(pipeline.transformer, original_params)
-
-        # validate lora_loaded map is updated
-        assert "adapter0" not in pipeline.lora_loaded
-
-    def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture):
-        pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
-        to_out = pipeline.transformer.blocks[0].attn.to_out
-        original_weight = to_out.weight.detach().clone()
-        base_key = "blocks.0.attn.to_out.0"
-        lora_a = torch.randn(RANK, HEAD_DIM)
-        lora_b = torch.randn(HEAD_DIM, RANK)
-        alpha = torch.tensor(6.0)
-        state_dict = {
-            f"{base_key}.lora_A.weight": lora_a,
-            f"{base_key}.lora_B.weight": lora_b,
-            f"{base_key}.alpha": alpha,
-        }
-        mocker.patch(
-            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
-            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
-        )
-
-        pipeline.load_lora_weights(state_dict, "adapter0")
-
-        expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK)
-        assert_close(to_out.weight, expected_weight)
-
-        pipeline.unload_lora_weights("adapter0")
-        assert_close(to_out.weight, original_weight)
-
-
-# ======================================================
-# Test WanLoraLoaderMixin
-# ======================================================
-class DummyWanPipeline(nn.Module, WanLoraLoaderMixin):
-    def __init__(self, num_layers: int, d_model: int):
-        super().__init__()
-        self.has_transformer_2 = True
-        self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True)
-        self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True)
-
-
-def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
-    pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM)
-    original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
-    original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()}
-    high_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
-    low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2)
-
-    mocker.patch(
-        "vllm_omni.diffusion.lora.loader._load_lora_state_dict",
-        side_effect=[high_state_dict, low_state_dict],
-    )
-
-    pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0")
-
-    assert_params_not_equal(pipeline.transformer, original_high)
-    assert_params_not_equal(pipeline.transformer_2, original_low)
-    assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"}
-
-    pipeline.unload_lora_weights("adapter0")
-
-    assert_params_equal(pipeline.transformer, original_high)
-    assert_params_equal(pipeline.transformer_2, original_low)
-    assert "adapter0" not in pipeline.lora_loaded

tests/diffusion/models/lingbot_video/test_request_utils.py

@@ -348,53 +348,3 @@ def test_normalize_lingbot_request_model_size_overrides_sampling_dimensions():
         num_frames=81,
     )
     assert (config.height, config.width) == (480, 480)
-
-
-def test_normalize_lingbot_request_reads_size_from_extra_args():
-    """Regression: /v1/videos delivers `extra_params` via sampling.extra_args.
-
-    The nightly TI2V case sends `extra_params={"size": "320x192"}` with no
-    height/width form fields; the size used to be dropped and the video came
-    out at the 480x480 default (build 2952, LingBot Video MoE Function Test).
-    """
-    config = _normalize(
-        {
-            "prompt": "motion",
-            "modalities": ["video"],
-            "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
-        },
-        num_frames=9,
-        extra_args={"size": "320x192"},
-    )
-    assert config.mode is LingBotGenerationMode.TI2V
-    assert (config.height, config.width) == (192, 320)
-
-
-def test_normalize_lingbot_request_extra_args_size_overrides_sampling_dimensions():
-    config = _normalize(
-        {"prompt": "motion", "modalities": ["video"]},
-        width=480,
-        height=480,
-        num_frames=81,
-        extra_args={"size": "320x192"},
-    )
-    assert (config.height, config.width) == (192, 320)
-
-
-def test_normalize_lingbot_request_reads_width_height_from_extra_args():
-    # Parity with the pre-#5311 pipeline, which honored extra_args height/width.
-    config = _normalize(
-        {"prompt": "motion", "modalities": ["video"]},
-        num_frames=81,
-        extra_args={"width": 320, "height": 192},
-    )
-    assert (config.height, config.width) == (192, 320)
-
-
-def test_resolve_lingbot_output_dimensions_reads_extra_fields():
-    assert resolve_lingbot_output_dimensions(extra_fields={"size": "320x192"}) == (320, 192)
-    assert resolve_lingbot_output_dimensions(
-        sampling_width=480,
-        sampling_height=480,
-        extra_fields={"resolution": "720p", "ratio": "16:9"},
-    ) == (736, 1280)

tests/diffusion/quantization/test_quantization_quality.py

@@ -515,14 +515,8 @@ def test_quantization_quality(config: QualityTestConfig):
 
     generate_fn = _generate_video if config.task == "t2v" else _generate_image
 
-    # Run both arms eager: LPIPS must measure quantization only. With compile
-    # on, a compiler failure in one arm (e.g. inductor's CantSplit on fp8 FLUX)
-    # silently drops that arm to eager while the other stays compiled, and the
-    # metric then includes compile-state differences (fp8_ltx2 rose from ~0.09
-    # to 0.1291 that way in build 2954). Mirrors
-    # vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py.
     # --- BF16 baseline ---
-    bl_kwargs: dict = {"model": config.baseline_ref(), "enforce_eager": True}
+    bl_kwargs: dict = {"model": config.baseline_ref()}
     if config.enable_cpu_offload:
         bl_kwargs["enable_cpu_offload"] = True
     omni_bl = Omni(**bl_kwargs)
@@ -534,7 +528,7 @@ def test_quantization_quality(config: QualityTestConfig):
 
     # --- Quantized ---
     quantization = config.quantization_ref()
-    qt_kwargs: dict = {"model": config.quantized_ref(), "enforce_eager": True}
+    qt_kwargs: dict = {"model": config.quantized_ref()}
     if config.enable_cpu_offload:
         qt_kwargs["enable_cpu_offload"] = True
     if quantization is None:

tests/diffusion/test_compile.py

@@ -74,37 +74,3 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
         regionally_compile(model, dynamic=True)
 
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
-
-
-def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch):
-    """cache-dit matches blocks via inspect.signature(block.forward).
-
-    Whatever regionally_compile installs as the block forward must stay
-    signature-transparent: parameter names and the return annotation drive
-    cache-dit's ForwardPattern match (build 2954, Multi-GPU Layered job
-    failed when a bare *args/**kwargs wrapper hid them; torch.compile's own
-    wrapper preserves the signature).
-    """
-    import inspect
-
-    import torch
-
-    class _SignatureBlock(nn.Module):
-        def forward(self, hidden_states, encoder_hidden_states=None) -> "torch.Tensor":
-            return hidden_states
-
-    class _Model(nn.Module):
-        _repeated_blocks = ["_SignatureBlock"]
-
-        def __init__(self) -> None:
-            super().__init__()
-            self.blocks = nn.ModuleList([_SignatureBlock()])
-
-    model = _Model()
-    monkeypatch.setattr(compile_module.torch, "compile", lambda fn, *a, **k: fn)
-    regionally_compile(model)
-
-    sig = inspect.signature(model.blocks[0].forward)
-    assert set(sig.parameters.keys()) == {"hidden_states", "encoder_hidden_states"}
-    assert "torch.Tensor" in str(sig.return_annotation)
-    assert model.blocks[0].forward("x") == "x"

tests/diffusion/test_inductor_divisibility_patch.py

@@ -1,83 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for the inductor factorable-divisibility patch.
-
-torch 2.13 proves symbolic divisibility with torch's own unevaluated Mod,
-so ``15360*s31 + 15360*s87`` is no longer provably divisible by
-``s31 + s87`` and inductor raises CantSplit on FLUX's fp8 graph (nightly
-builds 2953/2954). The patch adds an exact-polynomial-division proof that
-is sound for every symbol assignment; on torch versions whose original
-check already proves the case (<= 2.11, via sympy.Mod), the patch
-self-extinguishes.
-"""
-
-import pytest
-import sympy
-from torch._inductor.sizevars import SizeVarAllocator
-
-import vllm_omni.patch as patch_module
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-_S1, _S2 = sympy.symbols("s31 s87", positive=True, integer=True)
-
-
-@pytest.mark.parametrize(
-    ("numerator", "denominator", "expected"),
-    [
-        (15360 * _S1 + 15360 * _S2, _S1 + _S2, True),  # the CantSplit expression
-        (7 * _S1 + 7 * _S2, _S1 + _S2, True),
-        (_S1 * _S2, _S2, True),
-        (_S1 * _S2 + _S2, _S2, True),  # s2*(s1 + 1)
-        (3 * _S1, 2 * _S1, False),
-        (_S1 + 1, _S1, False),
-        (15360 * _S1 + 15359 * _S2, _S1 + _S2, False),
-        (sympy.Integer(12), sympy.Integer(4), True),
-        (sympy.Integer(12), sympy.Integer(5), False),
-    ],
-)
-def test_provably_factorable_multiple(numerator, denominator, expected):
-    assert patch_module._provably_factorable_multiple(numerator, denominator) is expected
-
-
-def test_non_polynomial_atoms_never_reach_cancel():
-    from torch.utils._sympy.functions import FloorDiv, ModularIndexing
-
-    assert patch_module._provably_factorable_multiple(FloorDiv(_S1, 2) * 4, _S1) is False
-    assert patch_module._provably_factorable_multiple(ModularIndexing(_S1, 1, 8) * _S2, _S2) is False
-
-
-def test_divisibility_holds_after_patch_regardless_of_torch_version():
-    """The end contract on any torch: the CantSplit expression proves divisible.
-
-    On torch <= 2.11 the original method proves it (patch self-extinguished);
-    on torch >= 2.13 the installed wrapper proves it via exact division.
-    """
-    sv = SizeVarAllocator()
-    method = SizeVarAllocator.statically_known_multiple_of
-    patched = getattr(method, "_vllm_omni_factorable_divisibility", False)
-    assert method(sv, 15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
-    # Negative control must stay non-divisible either way.
-    assert method(sv, 3 * _S1, 2 * _S1) is False
-    if patched:
-        # Wrapper installed: the original alone must have failed the probe.
-        assert patch_module._provably_factorable_multiple(15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
-
-
-def test_patch_installs_when_original_cannot_prove(monkeypatch):
-    """Simulate torch 2.13's unevaluated-Mod prover: the wrapper must install."""
-
-    def weak(self, numerator, denominator):
-        return False
-
-    monkeypatch.setattr(SizeVarAllocator, "statically_known_multiple_of", weak)
-    patch_module._patch_inductor_factorable_divisibility()
-    method = SizeVarAllocator.statically_known_multiple_of
-    assert getattr(method, "_vllm_omni_factorable_divisibility", False) is True
-
-    sv = SizeVarAllocator()
-    assert method(sv, 15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
-    assert method(sv, 3 * _S1, 2 * _S1) is False
-    # Idempotency: a second install call must not wrap the wrapper.
-    patch_module._patch_inductor_factorable_divisibility()
-    assert SizeVarAllocator.statically_known_multiple_of is method

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -33,19 +33,16 @@ BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
 # Generated with seed=52, num_inference_steps=15,
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
-# Re-goldened from the vLLM-0.27-aligned output (nightly builds 2953/2954
-# produced bit-identical pixels across two different torch/base images; only
-# (512, 336) moved beyond tolerance vs the main-branch goldens from #5898).
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (157, 172, 217)},
+    {"position": (100, 100), "rgb": (156, 172, 217)},
     {"position": (400, 50), "rgb": (105, 144, 217)},
     {"position": (700, 100), "rgb": (118, 159, 232)},
-    {"position": (150, 400), "rgb": (185, 21, 55)},
-    {"position": (512, 336), "rgb": (205, 215, 191)},
-    {"position": (700, 400), "rgb": (193, 11, 48)},
-    {"position": (100, 600), "rgb": (103, 14, 25)},
-    {"position": (400, 600), "rgb": (159, 28, 48)},
-    {"position": (700, 600), "rgb": (101, 87, 94)},
+    {"position": (150, 400), "rgb": (180, 22, 52)},
+    {"position": (512, 336), "rgb": (221, 211, 194)},
+    {"position": (700, 400), "rgb": (192, 10, 46)},
+    {"position": (100, 600), "rgb": (102, 12, 22)},
+    {"position": (400, 600), "rgb": (161, 28, 47)},
+    {"position": (700, 600), "rgb": (100, 87, 94)},
     {"position": (256, 256), "rgb": (181, 201, 221)},
 ]
 
@@ -144,24 +141,13 @@ def _validate_pixels(
     Raises:
         AssertionError: If any pixel differs beyond tolerance.
     """
-    mismatches = []
-    probes = []
     for ref in reference_pixels:
         x, y = ref["position"]
         expected = ref["rgb"]
         actual = image.getpixel((x, y))[:3]
-        probes.append(f'    {{"position": ({x}, {y}), "rgb": {actual}}},')
-        if not all(abs(a - e) <= tolerance for a, e in zip(actual, expected)):
-            mismatches.append(f"({x}, {y}): expected {expected}, got {actual}")
-    # Report every probe, not just the first mismatch: CI does not upload the
-    # generated image, so this failure message is the only place the full
-    # actual-value table exists when the goldens need to be regenerated.
-    assert not mismatches, (
-        f"Pixel mismatch at {len(mismatches)} probe(s) (tolerance ±{tolerance}):\n  "
-        + "\n  ".join(mismatches)
-        + "\nActual values at all probes (paste as the new reference table if re-goldening):\n"
-        + "\n".join(probes)
-    )
+        assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
+            f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
+        )
 
 
 def _generate_bagel_img2img(

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -1094,12 +1094,6 @@ class _HashableRequest(SimpleNamespace):
     # request; real Requests initialise it to 0 (vllm/v1/request.py).
     num_in_flight_tokens = 0
 
-    # vLLM 0.27 (a0c092ee72): the stale-output counter that replaced
-    # async_tokens_to_discard. update_from_output reads it for every scheduled
-    # request, and real Requests initialise it to 0 (vllm/v1/request.py), so the
-    # double needs it too.
-    num_stale_output_tokens = 0
-
     def __hash__(self):
         return hash(self.request_id)
 

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -42,31 +42,10 @@ from tests.e2e.accuracy.helpers import (
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer, OpenAIClientHandler
-from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import (
-    CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS,
-)
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
-def _set_matched_attention_backend(pipe: DiffusionPipeline) -> None:
-    """Walk the same backend chain the omni server's diffusers adapter walks.
-
-    The server side (--diffusion-load-format diffusers) resolves its attention
-    backend through the adapter's preference chain, skipping backends that are
-    unavailable on the image (e.g. the FA3 hub kernel has no build variant for
-    the image's torch — build 2953). The reference run must resolve to the
-    same backend or the similarity comparison measures kernel differences.
-    """
-    for backend in CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS:
-        try:
-            pipe.transformer.set_attention_backend(backend)
-            return
-        except Exception:
-            continue
-    pipe.transformer.set_attention_backend("native")
-
-
 PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
 NEGATIVE_PROMPT = "blurry, low quality"
 WIDTH = 512
@@ -256,7 +235,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
         ).to("cuda")
-        _set_matched_attention_backend(pipe)
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         _diffusers_dummy_run(pipe)
 
         latencies: list[float] = []

tests/e2e/accuracy/test_qwen_image.py

@@ -30,12 +30,6 @@ NUM_INFERENCE_STEPS = 20
 TRUE_CFG_SCALE = 4.0
 SEED = 42
 SSIM_THRESHOLD = 0.97
-# The FA3-hub-matched comparison clears 0.97 with margin. The SDPA-matched
-# fallback (hub kernel unavailable for the image's torch) landed at 0.9687 in
-# build 2954 — the attention math matches but the surrounding stacks (vLLM
-# compiled path vs eager diffusers) still differ slightly. Keep the strict
-# gate for the FA3 path and a slightly looser one for the fallback only.
-SSIM_THRESHOLD_SDPA_FALLBACK = 0.96
 PSNR_THRESHOLD = 30.0
 
 MODEL_2512_ID = "Qwen/Qwen-Image-2512"
@@ -62,44 +56,6 @@ MEAN_ABS_DIFF_THRESHOLD_2512 = 3e-2
 P99_ABS_DIFF_THRESHOLD_2512 = 4e-1
 
 
-_FLASH_ATTN3_HUB_AVAILABLE: bool | None = None
-
-
-def _flash_attn3_hub_available() -> bool:
-    """Whether kernels-community/flash-attn3 has a build for this torch/CUDA.
-
-    The hub repo only ships prebuilt variants for specific torch versions;
-    a base-image torch bump can outrun it (build 2952 failed with "Cannot
-    find a build variant for this system"). When it is unavailable, both
-    sides of the comparison fall back to torch SDPA together — the test
-    compares implementations, so the attention math must stay matched.
-    """
-    global _FLASH_ATTN3_HUB_AVAILABLE
-    if _FLASH_ATTN3_HUB_AVAILABLE is None:
-        try:
-            from kernels import get_kernel
-
-            get_kernel("kernels-community/flash-attn3")
-            _FLASH_ATTN3_HUB_AVAILABLE = True
-        except Exception as exc:
-            print(f"kernels-community/flash-attn3 unavailable ({exc}); using matched torch SDPA on both sides.")
-            _FLASH_ATTN3_HUB_AVAILABLE = False
-    return _FLASH_ATTN3_HUB_AVAILABLE
-
-
-def _omni_server_env() -> dict[str, str] | None:
-    if _flash_attn3_hub_available():
-        return None
-    return {"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA"}
-
-
-def _set_reference_attention_backend(pipe: DiffusionPipeline) -> None:
-    if _flash_attn3_hub_available():
-        pipe.transformer.set_attention_backend("_flash_3_hub")
-    else:
-        pipe.transformer.set_attention_backend("native")
-
-
 def _model_name() -> str:
     return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
 
@@ -122,7 +78,7 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
         "900",
         "--fa-deterministic",
     ]
-    with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server:
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
         response = requests.post(
             f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
             json={
@@ -158,7 +114,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
             trust_remote_code=True,
             local_files_only=_local_files_only(model),
         ).to("cuda")
-        _set_reference_attention_backend(pipe)
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         generator = torch.Generator(device="cuda").manual_seed(SEED)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=PROMPT,
@@ -184,7 +140,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
 
 def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
     server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
-    with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server:
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
         response = requests.post(
             f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
             json={
@@ -220,7 +176,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im
             trust_remote_code=True,
             local_files_only=_local_files_only(model),
         ).to("cuda")
-        _set_reference_attention_backend(pipe)
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         generator = torch.Generator(device="cuda").manual_seed(SEED_2512)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=PROMPT_2512,
@@ -259,7 +215,7 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
         diffusers_image=diffusers_output,
         width=WIDTH,
         height=HEIGHT,
-        ssim_threshold=SSIM_THRESHOLD if _flash_attn3_hub_available() else SSIM_THRESHOLD_SDPA_FALLBACK,
+        ssim_threshold=SSIM_THRESHOLD,
         psnr_threshold=PSNR_THRESHOLD,
     )
 

tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py

@@ -19,11 +19,38 @@ DEPLOY_CONFIG = modify_stage_config(
     get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
     updates={
         "base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
-        # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
-        # is left automatic so duplex matches the simplex deploy profiles.
+        # Cap per-stage KV so Thinker/Talker/Code2Wav share one GPU with
+        # max_sessions=2. Talker must stay within its 4096 context so the
+        # 0.5 GiB budget passes vLLM's min-KV check at init.
         "stages": {
-            1: {"max_model_len": 4096},
+            0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024},
+            1: {
+                "max_model_len": 4096,
+                "kv_cache_memory_bytes": 512 * 1024 * 1024,
+            },
+            2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
         },
+        # Platform overrides apply after ordinary stage settings and would
+        # otherwise reinstate the base CUDA 2 GiB Talker default.
+        "platforms": {
+            "cuda": {
+                "stages": [
+                    {
+                        "stage_id": 1,
+                        "kv_cache_memory_bytes": 512 * 1024 * 1024,
+                    }
+                ]
+            }
+        },
+    },
+)
+CORE_DEPLOY_CONFIG = modify_stage_config(
+    DEPLOY_CONFIG,
+    updates={
+        "stages": {
+            0: {"enforce_eager": True},
+            1: {"enforce_eager": True},
+        }
     },
 )
 ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
@@ -44,6 +71,17 @@ SERVER_PARAMS = [
         id="three-stage-single-gpu",
     )
 ]
+CORE_SERVER_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=CORE_DEPLOY_CONFIG,
+            use_stage_cli=False,
+            server_args=["--trust-remote-code"],
+        ),
+        id="three-stage-single-gpu",
+    )
+]
 
 
 def validated_wav(path: Path, expected_sha256: str) -> Path:
@@ -74,11 +112,9 @@ def validated_soft_interrupt_wav() -> Path:
 
 
 def resolve_ref_audio(model_prefix: str) -> Path:
-    # ``MODEL`` may be a local checkpoint directory instead of a Hub repo id
-    # (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only
-    # fall back to the Hub cache when no such directory exists on disk.
-    model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
-    if not model_root.is_dir():
+    if model_prefix:
+        model_root = Path(model_prefix) / MODEL
+    else:
         model_root = Path(snapshot_download(MODEL, local_files_only=True))
     ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
     if not ref_audio.is_file():

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -14,15 +14,47 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
 MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
-BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
 
 
 def _get_diffusion_feature_cases(model: str):
     return [
         pytest.param(
-            OmniServerParams(model=model),
-            id="base",
-            marks=BASE_FEATURE_MARKS,
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--tensor-parallel-size",
+                    "2",
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="parallel_tp_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--cfg-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_cfg_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--tensor-parallel-size",
+                    "2",
+                    "--enable-cpu-offload",
+                    "--cache-backend",
+                    "cache_dit",
+                ],
+            ),
+            id="parallel_002",
         ),
     ]
 

tests/e2e/online_serving/test_minicpmo_4_5.py

@@ -1,7 +1,10 @@
-"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
+"""
+E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
 
-Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
-Talker, and Code2Wav stages.
+MiniCPM-o 4.5 is a three-stage pipeline: Stage 0 Thinker (text), Stage 1 Talker
+(TTS codec), Stage 2 Code2Wav (vocoder). The talker -> Code2Wav bridge supports
+both async_chunk (streaming) and full-payload (``--no-async-chunk``) modes, so the
+tests below exercise the sync path as well.
 """
 
 import os
@@ -11,17 +14,12 @@ import pytest
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
 from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.stage_config import get_deploy_config_path
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 _MODEL = "openbmb/MiniCPM-o-4_5"
-_CI_DEPLOY = modify_stage_config(
-    get_deploy_config_path("minicpmo_4_5.yaml"),
-    updates={
-        "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}}
-    },
-)
+_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 
 test_params = [
     pytest.param(
@@ -29,10 +27,12 @@ test_params = [
             model=_MODEL,
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
-            server_args=["--trust-remote-code", "--async-chunk"],
+            server_args=[
+                "--trust-remote-code",
+            ],
         ),
-        id="async_chunk",
-    ),
+        id="default",
+    )
 ]
 
 
@@ -98,7 +98,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
 def test_text_to_audio_001(omni_server, openai_client) -> None:
     """
     Test text-only input generating text + audio output via OpenAI API.
-    This exercises Talker TTS region detection and the Code2Wav stage.
+    This exercises the talker TTS region detection and in-process token2wav vocoder.
     Deploy Setting: default single GPU
     Input Modal: text
     Output Modal: text + audio
@@ -132,21 +132,13 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
     messages = dummy_messages_from_mix_data(
         system_prompt=get_system_prompt(),
         audio_data_url=audio_data_url,
-        content_text=get_prompt(),
+        content_text=get_prompt("mix"),
     )
 
     request_config = {
         "model": omni_server.model,
         "messages": messages,
         "stream": True,
-        "key_words": {"text": ["Beijing"]},
-        "modalities": ["text", "audio"],
-        "extra_body": {
-            "chat_template_kwargs": {
-                "use_tts_template": True,
-                "enable_thinking": False,
-            }
-        },
     }
 
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_minicpmo_4_5_duplex.py

@@ -11,6 +11,7 @@ import pytest
 import websockets
 
 from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+    CORE_SERVER_PARAMS,
     SERVER_PARAMS,
     demo_args,
     multi_session_args,
@@ -106,7 +107,7 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[
 
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True)
 def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
     ref_audio = resolve_ref_audio(model_prefix)
     events = asyncio.run(
@@ -133,12 +134,8 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
         output_dir=tmp_path / "single_session",
     )
     args.turns = 2
-    # Every turn replays the same active-speech window as the first one. The default
-    # shorter follow-up window is a different mid-utterance slice, which the native
-    # duplex model may legitimately answer with "listen" instead of a response.
-    args.turn_duration_ms = [args.first_turn_ms] * args.turns
     result = asyncio.run(run_demo(args))
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
     assert result["audio_delta_count"] > 0
     assert result["done_count"] == 2
     assert result["error_count"] == 0
@@ -163,7 +160,7 @@ def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str,
             )
         )
     )
-    assert result["ok"] is True
+    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
     assert result["session_count"] == 2
     assert result["resume"]["ok"] is True
     assert result["takeover"]["ok"] is True

tests/e2e/online_serving/test_minicpmo_4_5_expansion.py

@@ -2,9 +2,6 @@
 
 These cover modality combinations, separate Code2Wav behavior, sequential
 request isolation, longer audio output, and Chinese speech.
-
-Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and
-``async_chunk`` exercises ``--async-chunk``.
 """
 
 import os
@@ -23,27 +20,18 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 _MODEL = "openbmb/MiniCPM-o-4_5"
 _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 
-# Deploy YAML defaults to ``async_chunk: true``; keep sync and async as
-# separate parametrizations so both handoff paths are exercised.
 test_params = [
     pytest.param(
         OmniServerParams(
             model=_MODEL,
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
-            server_args=["--trust-remote-code", "--no-async-chunk"],
+            server_args=[
+                "--trust-remote-code",
+            ],
         ),
         id="default",
-    ),
-    pytest.param(
-        OmniServerParams(
-            model=_MODEL,
-            stage_config_path=_CI_DEPLOY,
-            use_stage_cli=False,
-            server_args=["--trust-remote-code", "--async-chunk"],
-        ),
-        id="async_chunk",
-    ),
+    )
 ]
 
 

tests/e2e/online_serving/test_voxtral_tts_expansion.py

@@ -118,17 +118,6 @@ def test_speech_instructions(omni_server, openai_client) -> None:
                 "response_format": "wav",
                 "timeout": 120.0,
                 "instructions": instruction,
-                # Re-grade a below-threshold result with whisper large-v3 before
-                # failing, so a miss is attributable rather than ambiguous. Same
-                # opt-in as test_higgs_audio_v3.py.
-                #
-                # Note: for this input it does NOT rescue the check — small and
-                # large-v3 both transcribe "The boy was there when sun rose."
-                # (sim=0.87 each) against an input of "...when the sun rose.",
-                # i.e. the TTS really does drop the article. That is the value
-                # here: two independent ASR models agreeing rules out a grader
-                # flake and points at the model.
-                "transcript_escalation_model": "large-v3",
             }
         )
 

tests/engine/test_stage_engine_startup_cache_env.py

@@ -1,47 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for per-replica compile-cache isolation at spawn time.
-
-Replicas of the same stage compile the same model concurrently, and vLLM
-0.27's AOT compile (TritonBundler) races on the shared VLLM_CACHE_ROOT —
-a replica then dies with 'Cubin file saved by TritonBundler not found'
-(build 2954, Omni Perf Multi-Replica). Replica 0 must keep the default
-shared cache; additional replicas must spawn with an isolated cache dir,
-and the parent env must be restored afterwards.
-"""
-
-import os
-
-import pytest
-
-from vllm_omni.engine.stage_engine_startup import scoped_spawn_device_env
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_replica_zero_keeps_shared_cache_root(monkeypatch):
-    monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False)
-    with scoped_spawn_device_env(None, None, stage_id=1, replica_id=0):
-        assert "VLLM_CACHE_ROOT" not in os.environ
-    assert "VLLM_CACHE_ROOT" not in os.environ
-
-
-def test_additional_replica_gets_isolated_cache_root(monkeypatch):
-    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/base-cache")
-    with scoped_spawn_device_env(None, None, stage_id=2, replica_id=1):
-        assert os.environ["VLLM_CACHE_ROOT"] == os.path.join("/tmp/base-cache", "omni_stage2_replica1")
-    assert os.environ["VLLM_CACHE_ROOT"] == "/tmp/base-cache"
-
-
-def test_isolated_cache_root_defaults_without_env(monkeypatch):
-    monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False)
-    with scoped_spawn_device_env(None, None, stage_id=0, replica_id=3):
-        assert os.environ["VLLM_CACHE_ROOT"].endswith("omni_stage0_replica3")
-        assert os.environ["VLLM_CACHE_ROOT"].startswith(os.path.expanduser("~/.cache/vllm"))
-    assert "VLLM_CACHE_ROOT" not in os.environ
-
-
-def test_missing_ids_leave_env_untouched(monkeypatch):
-    monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False)
-    with scoped_spawn_device_env(None, None):
-        assert "VLLM_CACHE_ROOT" not in os.environ

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -284,37 +284,6 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
 
 
-def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
-    """Ensure startup distilled LoRA options reach the online diffusion stage."""
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="command")
-    OmniServeCommand().subparser_init(subparsers)
-
-    args = parser.parse_args(
-        [
-            "serve",
-            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-            "--omni",
-            "--lora-backend",
-            "distill",
-            "--lora-path",
-            "/models/high.safetensors",
-            "/models/low.safetensors",
-        ]
-    )
-
-    explicit_kwargs = args.get_explicit_kwargs_dict()
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
-    engine_args = stage_cfg["engine_args"]
-
-    assert explicit_kwargs["lora_backend"] == "distill"
-    assert engine_args["lora_backend"] == "distill"
-    assert engine_args["lora_path"] == [
-        "/models/high.safetensors",
-        "/models/low.safetensors",
-    ]
-
-
 def test_serve_cli_forwards_distributed_offload_residency():
     """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
     parser = TrackingArgumentParser()

tests/examples/offline_inference/test_image_task_prompts.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+from PIL import Image
+
+from examples.offline_inference.image_to_video.image_to_video import build_image_to_video_prompt
+from examples.offline_inference.text_to_image.text_to_image import build_text_to_image_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("negative_prompt", "expected_negative_prompt"),
+    [(None, None), ("", ""), ("blurry", "blurry")],
+)
+def test_text_to_image_builds_canonical_prompt(
+    negative_prompt: str | None,
+    expected_negative_prompt: str | None,
+) -> None:
+    result = build_text_to_image_prompt("a red fox", negative_prompt)
+
+    assert result["prompt"] == "a red fox"
+    assert result["modalities"] == ["image"]
+    if expected_negative_prompt is None:
+        assert "negative_prompt" not in result
+    else:
+        assert result["negative_prompt"] == expected_negative_prompt
+
+
+def test_image_to_video_builds_canonical_prompt() -> None:
+    image = Image.new("RGB", (32, 16), "red")
+
+    result = build_image_to_video_prompt(
+        prompt="the fox turns toward the camera",
+        negative_prompt="flicker",
+        media_inputs={"image": image},
+    )
+
+    assert result == {
+        "prompt": "the fox turns toward the camera",
+        "modalities": ["video"],
+        "multi_modal_data": {"image": image},
+        "negative_prompt": "flicker",
+    }

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -285,17 +285,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         captured["mm_cache_keys"] = kwargs["mm_hashes"]
         return MagicMock(name="mm_kwargs"), {"video": [[]], "audio": [[]]}
 
-    fake_self = SimpleNamespace(
-        cache=cache,
-        info=SimpleNamespace(
-            model_id="m",
-            ctx=SimpleNamespace(
-                get_mm_config=lambda: SimpleNamespace(
-                    mm_hasher_algorithm="sha256",
-                )
-            ),
-        ),
-    )
+    fake_self = SimpleNamespace(cache=cache, info=SimpleNamespace(model_id="m"))
     fake_self._get_hf_mm_data = MagicMock(return_value=({}, {}))
     # Bound methods called via self inside _cached_apply_hf_processor.
     fake_self._get_audio_in_video_pairs = lambda mm_data_items, hf_processor_mm_kwargs: (

tests/model_executor/models/qwen3_omni/test_qwen3_omni_forward_contract.py

@@ -1,175 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for the Qwen3-Omni thinker forward return contract.
-
-Background
-----------
-Since omni architectures unconditionally override upstream's registry entries,
-``Qwen3OmniMoeForConditionalGeneration`` (and the standalone thinker) are also
-constructed by plain, non-staged vLLM runs — e.g. ``vllm serve
-Qwen/Qwen3-Omni-30B-A3B-Instruct`` in the nightly vLLM-text perf job. Stock
-vLLM's ``GPUModelRunner`` indexes the forward output directly
-(``hidden_states[logit_indices]``), so a thinker forward that always returns a
-``(hidden_states, captured_hidden_states)`` tuple kills the engine during
-``profile_run`` with ``TypeError: only integer tensors of a single element can
-be converted to an index`` (Buildkite build 2952).
-
-Contract under test:
-
-- No capture requested → the forward returns a **bare tensor** (stock-runner
-  compatible).
-- Capture requested (``capture_layer_indices`` + ``return_hidden_states``,
-  staged runs with a talker downstream) → the forward returns the
-  ``(hidden_states, captured_layer_dict)`` tuple omni's runner consumes via
-  ``make_omni_output``.
-- The combined model requests capture only in staged runs
-  (``model_config.model_stage`` injected), never in plain vLLM runs.
-- ``make_omni_output`` tolerates both return shapes.
-
-Pure shape/contract tests: CPU-only, no weights, no GPU.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from torch import nn
-from vllm.sequence import IntermediateTensors
-
-from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import (
-    Qwen3OmniMoeForConditionalGeneration,
-)
-from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
-    Qwen3OmniMoeThinkerForConditionalGeneration,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_HIDDEN = 8
-_TOKENS = 4
-
-
-def _hidden_states() -> torch.Tensor:
-    return torch.zeros(_TOKENS, _HIDDEN)
-
-
-class _InnerModelStub:
-    """Stands in for ``Qwen3MoeLLMModel``: always returns the internal 2-tuple."""
-
-    def __init__(self, output=None):
-        self.calls: list[dict] = []
-        self._output = output
-
-    def __call__(self, input_ids, positions, intermediate_tensors, **kwargs):
-        self.calls.append(kwargs)
-        if self._output is not None:
-            return self._output
-        if kwargs.get("return_hidden_states"):
-            return _hidden_states(), {"hidden_states": {"layers": {0: _hidden_states()}}}
-        return _hidden_states(), None
-
-
-def _make_thinker(inner_output=None) -> tuple[Qwen3OmniMoeThinkerForConditionalGeneration, _InnerModelStub]:
-    thinker = object.__new__(Qwen3OmniMoeThinkerForConditionalGeneration)
-    nn.Module.__init__(thinker)
-    inner = _InnerModelStub(output=inner_output)
-    thinker.use_deepstack = False
-    thinker.language_model = SimpleNamespace(model=inner)
-    return thinker, inner
-
-
-def _forward_args():
-    return {
-        "input_ids": torch.zeros(_TOKENS, dtype=torch.long),
-        "positions": torch.arange(_TOKENS),
-    }
-
-
-def test_thinker_forward_returns_bare_tensor_without_capture():
-    thinker, _ = _make_thinker()
-    out = thinker.forward(**_forward_args())
-    assert isinstance(out, torch.Tensor)
-    assert out.shape == (_TOKENS, _HIDDEN)
-
-
-def test_thinker_forward_returns_tuple_when_capturing():
-    thinker, _ = _make_thinker()
-    out = thinker.forward(
-        **_forward_args(),
-        capture_layer_indices=[0, 2],
-        return_hidden_states=True,
-    )
-    assert isinstance(out, tuple)
-    hidden, captured = out
-    assert isinstance(hidden, torch.Tensor)
-    assert "hidden_states" in captured
-
-
-def test_thinker_forward_passes_through_intermediate_tensors():
-    it = IntermediateTensors({"hidden_states": _hidden_states(), "residual": _hidden_states()})
-    thinker, _ = _make_thinker(inner_output=it)
-    out = thinker.forward(**_forward_args())
-    assert out is it
-
-
-class _ThinkerStub(nn.Module):
-    """Records capture kwargs; mirrors the fixed thinker return contract."""
-
-    def __init__(self):
-        super().__init__()
-        self.received: dict = {}
-
-    def forward(self, **kwargs):
-        self.received = kwargs
-        if kwargs.get("return_hidden_states"):
-            return _hidden_states(), {"hidden_states": {"layers": {}}}
-        return _hidden_states()
-
-
-def _make_combined(*, is_staged_run: bool, accept_hidden_layer=6) -> Qwen3OmniMoeForConditionalGeneration:
-    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
-    nn.Module.__init__(model)
-    model.model_stage = "thinker"
-    model.is_staged_run = is_staged_run
-    model.talker_config = SimpleNamespace(accept_hidden_layer=accept_hidden_layer)
-    model.thinker = _ThinkerStub()
-    return model
-
-
-def test_combined_forward_plain_run_skips_capture_and_returns_tensor():
-    model = _make_combined(is_staged_run=False)
-    out = model.forward(**_forward_args())
-    assert "capture_layer_indices" not in model.thinker.received
-    assert isinstance(out, torch.Tensor)
-
-
-def test_combined_forward_staged_run_captures_and_returns_tuple():
-    model = _make_combined(is_staged_run=True)
-    out = model.forward(**_forward_args())
-    assert model.thinker.received["capture_layer_indices"] == [0, 6]
-    assert model.thinker.received["return_hidden_states"] is True
-    assert isinstance(out, tuple)
-
-
-def test_combined_forward_staged_run_without_accept_layer_returns_tensor():
-    model = _make_combined(is_staged_run=True, accept_hidden_layer=None)
-    out = model.forward(**_forward_args())
-    assert "capture_layer_indices" not in model.thinker.received
-    assert isinstance(out, torch.Tensor)
-
-
-def test_make_omni_output_accepts_bare_tensor():
-    model = _make_combined(is_staged_run=False)
-    out = model.make_omni_output(_hidden_states())
-    assert isinstance(out, OmniOutput)
-    assert out.text_hidden_states.shape == (_TOKENS, _HIDDEN)
-    assert out.multimodal_outputs == {}
-
-
-def test_make_omni_output_accepts_capture_tuple():
-    model = _make_combined(is_staged_run=True)
-    captured = {"hidden_states": {"layers": {}}}
-    out = model.make_omni_output((_hidden_states(), captured))
-    assert isinstance(out, OmniOutput)
-    assert out.multimodal_outputs == captured

tests/model_extras/test_model_extras.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import pytest
 from PIL import Image
 
@@ -135,67 +137,6 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
     assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_lingbot_text_to_image_prompt_builder_preserves_empty_negative_prompt() -> None:
-    assert build_text_to_image_prompt(
-        "LingBotVideoPipeline",
-        prompt="a red fox in fresh snow",
-        negative_prompt="",
-        height=192,
-        width=320,
-    ) == {
-        "prompt": "a red fox in fresh snow",
-        "modalities": ["image"],
-        "negative_prompt": "",
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_lingbot_image_to_video_prompt_builder() -> None:
-    image = Image.new("RGB", (320, 192), "red")
-    result = build_image_to_video_prompt(
-        "LingBotVideoPipeline",
-        prompt="the fox looks toward the camera",
-        negative_prompt=None,
-        media_inputs={"image": image},
-        height=192,
-        width=320,
-        num_frames=9,
-    )
-    assert result == {
-        "prompt": "the fox looks toward the camera",
-        "modalities": ["video"],
-        "multi_modal_data": {"image": image},
-        "num_frames": 9,
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "media_inputs",
-    [
-        {},
-        {"image": "input.png"},
-        {"image": [Image.new("RGB", (16, 16))]},
-        {"image": Image.new("RGB", (16, 16)), "video": b"video"},
-    ],
-    ids=["missing", "path", "multiple", "extra-modality"],
-)
-def test_lingbot_image_to_video_prompt_builder_rejects_invalid_media(
-    media_inputs: dict[str, object],
-) -> None:
-    with pytest.raises(ValueError, match="exactly one PIL image"):
-        build_image_to_video_prompt(
-            "LingBotVideoPipeline",
-            prompt="move",
-            negative_prompt=None,
-            media_inputs=media_inputs,
-        )
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 @pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
@@ -342,8 +283,11 @@ def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> Non
 def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
     assert build_text_to_image_prompt(
         "MingImagePipeline",
-        prompt="Please draw a cute cat.",
-        negative_prompt="ugly, blurry",
+        {
+            "prompt": "Please draw a cute cat.",
+            "modalities": ["image"],
+            "negative_prompt": "ugly, blurry",
+        },
         height=512,
         width=768,
     ) == {
@@ -360,8 +304,7 @@ def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
     # pipeline's 1024x1024 default still applies.
     assert build_text_to_image_prompt(
         "MingImagePipeline",
-        prompt="Draw a poster.",
-        negative_prompt=None,
+        {"prompt": "Draw a poster.", "modalities": ["image"]},
     ) == {
         "prompt": "Draw a poster.",
         "modalities": ["image"],
@@ -391,28 +334,6 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
     assert "negative_prompt" not in result
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
-def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None:
-    assert build_text_to_image_prompt(
-        pipeline_name,
-        prompt="a red sports car at golden hour",
-        negative_prompt="blurry, distorted",
-        height=1024,
-        width=1024,
-    ) == {
-        "prompt": "a red sports car at golden hour",
-        "modalities": ["image"],
-        "negative_prompt": "blurry, distorted",
-    }
-    assert build_text_to_image_prompt(
-        pipeline_name,
-        prompt="a red sports car",
-        negative_prompt=None,
-    ) == {"prompt": "a red sports car", "modalities": ["image"]}
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_audiox_extra_registry_declares_request_and_response_params() -> None:
@@ -505,8 +426,11 @@ def test_unknown_pipeline_has_empty_extra_registry() -> None:
 def test_bagel_text_to_image_prompt_builder() -> None:
     assert build_text_to_image_prompt(
         "BagelPipeline",
-        prompt="a cat",
-        negative_prompt="blurry",
+        {
+            "prompt": "a cat",
+            "modalities": ["image"],
+            "negative_prompt": "blurry",
+        },
         height=512,
         width=768,
     ) == {
@@ -543,14 +467,17 @@ def test_bagel_image_to_image_prompt_builder() -> None:
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_unknown_pipeline_uses_default_text_to_image_prompt() -> None:
-    assert build_text_to_image_prompt(
-        "UnknownPipeline",
-        prompt="a cat",
-        negative_prompt=None,
-        height=512,
-        width=512,
-    ) == {"prompt": "a cat"}
+def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None:
+    canonical_prompt = {"prompt": "a cat", "modalities": ["image"]}
+    assert (
+        build_text_to_image_prompt(
+            "UnknownPipeline",
+            canonical_prompt,
+            height=512,
+            width=512,
+        )
+        is canonical_prompt
+    )
 
 
 @pytest.mark.core_model
@@ -569,12 +496,14 @@ def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
     }
 
 
-def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict:
+def _build_vace_prompt(media_inputs: dict[str, Any], *, num_frames: int = 5) -> dict:
     return build_image_to_video_prompt(
         "WanVACEPipeline",
-        prompt="a bird flying",
-        negative_prompt=None,
-        media_inputs=media_inputs,
+        {
+            "prompt": "a bird flying",
+            "modalities": ["video"],
+            "multi_modal_data": media_inputs,
+        },
         height=16,
         width=320,
         num_frames=num_frames,
@@ -597,7 +526,7 @@ def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5)
     ],
     ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"],
 )
-def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None:
+def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, Any]) -> None:
     result = _build_vace_prompt(media_inputs)
     mmd = result["multi_modal_data"]
     if "reference_images" in mmd:
@@ -616,7 +545,7 @@ def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) ->
         ({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"),
     ],
 )
-def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None:
+def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, Any], message: str) -> None:
     with pytest.raises(ValueError, match=message):
         _build_vace_prompt(media_inputs)
 
@@ -664,8 +593,7 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N
 
     wrapper_prompt = build_text_to_image_prompt(
         "MammothModa2ForConditionalGeneration",
-        prompt="a cat",
-        negative_prompt=None,
+        {"prompt": "a cat", "modalities": ["image"]},
         height=256,
         width=256,
     )
@@ -679,8 +607,11 @@ def test_mammothmoda2_text_to_image_prompt_builder() -> None:
     # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
     assert build_text_to_image_prompt(
         "MammothModa2DiTPipeline",
-        prompt="a cat",
-        negative_prompt="blurry",
+        {
+            "prompt": "a cat",
+            "modalities": ["image"],
+            "negative_prompt": "blurry",
+        },
         height=512,
         width=768,
     ) == {

tests/model_tests/diffusion/diff_model_builders.py

@@ -37,19 +37,6 @@ def tiny_ltx2_builder() -> str:
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
 
 
-def _shrink_flux_clip_text_encoder(config: dict) -> dict:
-    config["num_hidden_layers"] = 2
-    return config
-
-
-def _shrink_flux_t5_text_encoder(config: dict) -> dict:
-    config["num_layers"] = 2
-    config["num_decoder_layers"] = 2
-    config["d_ff"] = 64
-    config["num_heads"] = 4
-    return config
-
-
 def tiny_qwen_image_builder() -> str:
     def shrink_text_encoder(config: dict) -> dict:
         config["num_hidden_layers"] = 2
@@ -70,24 +57,23 @@ def tiny_qwen_image_builder() -> str:
 
 
 def tiny_flux_builder() -> str:
+    def shrink_clip_text_encoder(config: dict) -> dict:
+        config["num_hidden_layers"] = 2
+        return config
+
+    def shrink_t5_text_encoder(config: dict) -> dict:
+        config["num_layers"] = 2
+        config["num_decoder_layers"] = 2
+        config["d_ff"] = 64
+        config["num_heads"] = 4
+        return config
+
     return build_tiny_from_configs(
         "FluxPipeline",
         "black-forest-labs/FLUX.1-schnell",
         transform={
-            "text_encoder": _shrink_flux_clip_text_encoder,
-            "text_encoder_2": _shrink_flux_t5_text_encoder,
-            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
-        },
-    )
-
-
-def tiny_flux_kontext_builder() -> str:
-    return build_tiny_from_configs(
-        "FluxKontextPipeline",
-        "black-forest-labs/FLUX.1-Kontext-dev",
-        transform={
-            "text_encoder": _shrink_flux_clip_text_encoder,
-            "text_encoder_2": _shrink_flux_t5_text_encoder,
+            "text_encoder": shrink_clip_text_encoder,
+            "text_encoder_2": shrink_t5_text_encoder,
             "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
         },
     )

tests/model_tests/diffusion/model_settings.py

@@ -59,16 +59,6 @@ DIFFUSION_TEST_SETTINGS = {
         builder=diff_model_builders.tiny_flux_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
     ),
-    "FluxKontextPipeline": DiffusionModelTestOpts(
-        model="black-forest-labs/FLUX.1-Kontext-dev",
-        builder=diff_model_builders.tiny_flux_kontext_builder,
-        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
-        extra_test_groups=[
-            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
-            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
-            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD, DiffusionAccs.CACHE_DIT],
-        ],
-    ),
     "Flux2Pipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.2-dev",
         builder=diff_model_builders.tiny_flux2_builder,

tests/model_tests/diffusion/test_alignment.py

@@ -48,6 +48,7 @@ EXCLUDED_MODELS = [
     "InternVLAA1Pipeline",
     "LongCatImageEditPipeline",
     "StableDiffusion3Pipeline",
+    "FluxKontextPipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",

tests/worker/test_worker_base.py

@@ -30,12 +30,9 @@ def _fake_memory_profiling(*, non_torch: int, torch_peak: int):
 
     @contextmanager
     def _mp(snapshot, weights_memory):  # noqa: ARG001 - signature parity only
-        # total_consumed mirrors upstream vllm.utils.mem_utils.MemoryProfilingResult
-        # (added by upstream 58b2012aa2) and is read by OmniGPUWorkerBase.
         yield SimpleNamespace(
             non_torch_increase=non_torch,
             torch_peak_increase=torch_peak,
-            total_consumed=torch_peak + non_torch,
         )
 
     return _mp

vllm_omni/benchmarks/patch/patch.py

@@ -12,7 +12,7 @@ import traceback
 import uuid
 import wave
 from collections.abc import Iterable
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Literal
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
@@ -1542,7 +1542,6 @@ async def benchmark(
     ready_check_timeout_sec: int = 600,
     ssl_context: ssl.SSLContext | bool | None = None,
     self_timed: bool = False,
-    probe_request_rate: float = 0.0,
 ):
     try:
         request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -1688,31 +1687,6 @@ async def benchmark(
         async with semaphore:
             return await request_func(request_func_input=request_func_input, session=session, pbar=pbar)
 
-    # Ported from upstream vLLM v0.27.0 (vllm/benchmarks/serve.py), which added
-    # the probe_request_rate background probe; this file is the patched copy of
-    # that module and tracks its upstream base.
-    probe_outputs: list[MixRequestFuncOutput] = []
-    probe_stop = asyncio.Event()
-
-    async def probe_loop():
-        probe_input = replace(
-            test_input,
-            prompt="Hi",
-            prompt_len=1,
-            output_len=1,
-            multi_modal_content=None,
-            chat_messages=None,
-        )
-        interval = 1 / probe_request_rate
-        while not probe_stop.is_set():
-            probe_outputs.append(await request_func(request_func_input=probe_input, session=session))
-            await asyncio.sleep(interval)
-
-    probe_task: asyncio.Task | None = None
-    if probe_request_rate > 0:
-        print(f"Probe request rate: {probe_request_rate} req/s")
-        probe_task = asyncio.create_task(probe_loop())
-
     benchmark_start_time = time.perf_counter()
     tasks: list[asyncio.Task] = []
 
@@ -1778,10 +1752,6 @@ async def benchmark(
         )
     outputs: list[MixRequestFuncOutput] = await asyncio.gather(*tasks)
 
-    if probe_task is not None:
-        probe_stop.set()
-        await probe_task
-
     if pbar is not None:
         pbar.close()
 
@@ -1849,19 +1819,11 @@ async def benchmark(
             "input_lens": [output.prompt_len for output in outputs],
             "errors": [output.error for output in outputs],
         }
-    # Plain-vLLM backends (e.g. the vLLM-text perf config) return upstream
-    # RequestFuncOutput objects without the Mix duplex fields; read them
-    # tolerantly or the whole benchmark result is discarded ("fallback to
-    # template", completed=0) after every request already succeeded.
-    duplex_request_metrics = [
-        metric for output in outputs for metric in (getattr(output, "duplex_request_metrics", None) or [])
-    ]
+    duplex_request_metrics = [metric for output in outputs for metric in (output.duplex_request_metrics or [])]
     if duplex_request_metrics:
         result["duplex_request_metrics"] = duplex_request_metrics
     duplex_session_metrics = [
-        session_metrics
-        for output in outputs
-        if (session_metrics := getattr(output, "duplex_session_metrics", None)) is not None
+        output.duplex_session_metrics for output in outputs if output.duplex_session_metrics is not None
     ]
     if duplex_session_metrics:
         result["duplex_session_metrics"] = duplex_session_metrics

vllm_omni/config/omni_config.py

@@ -563,9 +563,8 @@ class _DiffusionConfigProjection:
     diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_pipeline_cls: Any = None
-    lora_path: str | list[str] | None = None
+    lora_path: str | None = None
     lora_scale: float = 1.0
-    lora_backend: str = "peft"
     max_cpu_loras: int | None = None
     output_type: str = "pil"
     enable_cpu_offload: bool = False

vllm_omni/config/stage_config.py

@@ -385,9 +385,6 @@ class StageDeployConfig:
     # Diffusion model loading and adapter construction.
     model_class_name: str | None = None
     diffusion_load_format: str | None = None
-    lora_path: str | list[str] | None = None
-    lora_backend: str | None = None
-    lora_scale: float | None = None
     diffusers_load_kwargs: dict[str, Any] | None = None
     diffusers_call_kwargs: dict[str, Any] | None = None
     diffusion_quantization_config: str | None = None

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -325,18 +325,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # max(0, computed - in_flight), so a leaked counter silently
                 # freezes sliding-window block freeing.
                 request.num_in_flight_tokens -= num_tokens_scheduled
-            # vLLM 0.27 (a0c092ee72) removed the async_tokens_to_discard
-            # handling from the upstream scheduler and replaced it with the
-            # num_stale_output_tokens/is_stale mechanism. Omni's discard
-            # sites (segment stop, streaming-session replacement) record the
-            # in-flight share here; the delayed outputs are dropped below
-            # instead of decrementing num_output_placeholders (which the
-            # discard zeroed) and underflowing the upstream assert.
-            output_is_stale = False
-            if request is not None and request.num_stale_output_tokens > 0:
-                output_is_stale = True
-                request.num_stale_output_tokens -= num_tokens_scheduled
-                assert request.num_stale_output_tokens >= 0
             if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
                 # Skip requests that were recovered from KV load failure
                 continue
@@ -345,12 +333,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # request is aborted while the model is executing it (e.g.,
                 # in pipeline parallelism or async scheduling).
                 continue
-            if output_is_stale:
-                # Output of a step scheduled before the request's in-flight
-                # tokens were discarded (segment stop / session replacement).
-                # num_computed_tokens was rolled back at the discard site, so
-                # this output must not be appended or emitted.
-                continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
@@ -470,28 +452,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                             # connector. This update only resumes polling for the next segment.
                             self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
                     outstanding_async_tokens = request.num_output_placeholders
-                    # Always record the discard signal (0 when nothing is in
-                    # flight). Upstream a0c092ee72 removed the
-                    # `async_tokens_to_discard` default from `Request`; it
-                    # remains an omni-only signal set here on every segment
-                    # stop, so a stop with no outstanding placeholders
-                    # explicitly records 0.
-                    request.async_tokens_to_discard = outstanding_async_tokens
-                    # Seed the stale share in SCHEDULED-token units:
-                    # num_in_flight_tokens is exactly the unreported steps'
-                    # num_tokens_scheduled sum (settled per frame at the top
-                    # of this loop), and the drain subtracts each arriving
-                    # frame's num_tokens_scheduled — commensurable by
-                    # construction, so pre-discard frames drain to exactly
-                    # zero. Seeding from num_output_placeholders swallowed
-                    # valid new-segment frames or underflowed the drain
-                    # assert whenever placeholder counts diverged from
-                    # scheduled counts (spec drafts, in-flight prefill).
-                    if request.num_in_flight_tokens > 0:
-                        request.num_stale_output_tokens += request.num_in_flight_tokens
                     if outstanding_async_tokens > 0:
                         # Discard only outputs that are already in flight and
                         # roll back their optimistic computed-token accounting.
+                        request.async_tokens_to_discard = outstanding_async_tokens
                         request.num_computed_tokens -= outstanding_async_tokens
                         request.num_output_placeholders = 0
                     request.spec_token_ids = []
@@ -618,15 +582,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
         outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
-        # Seed the stale share in SCHEDULED-token units (see the segment-stop
-        # site in update_from_output): num_in_flight_tokens matches what each
-        # pre-replacement frame will drain, so the counter reaches exactly
-        # zero and the new segment's frames are never swallowed. This also
-        # covers an in-flight prefill chunk, which carries no placeholders
-        # but must still have its late output dropped.
-        in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
-        if in_flight_tokens > 0:
-            session.num_stale_output_tokens += in_flight_tokens
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous
             # segment's next token. Drop that late token instead of

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -374,18 +374,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # max(0, computed - in_flight), so a leaked counter silently
                 # freezes sliding-window block freeing.
                 request.num_in_flight_tokens -= num_tokens_scheduled
-            # vLLM 0.27 (a0c092ee72) removed the async_tokens_to_discard
-            # handling from the upstream scheduler and replaced it with the
-            # num_stale_output_tokens/is_stale mechanism. Omni's discard
-            # sites (segment stop, streaming-session replacement) record the
-            # in-flight share here; the delayed outputs are dropped below
-            # instead of decrementing num_output_placeholders (which the
-            # discard zeroed) and underflowing the upstream assert.
-            output_is_stale = False
-            if request is not None and request.num_stale_output_tokens > 0:
-                output_is_stale = True
-                request.num_stale_output_tokens -= num_tokens_scheduled
-                assert request.num_stale_output_tokens >= 0
             if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
                 # Skip requests that were recovered from KV load failure
                 continue
@@ -393,12 +381,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Request may already be finished (e.g., aborted during
                 # execution / pipeline parallelism / async scheduling).
                 continue
-            if output_is_stale:
-                # Output of a step scheduled before the request's in-flight
-                # tokens were discarded (segment stop / session replacement).
-                # num_computed_tokens was rolled back at the discard site, so
-                # this output must not be appended or emitted.
-                continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -94,17 +94,6 @@ class OmniSchedulerMixin:
             adapter.segment_finished_requests.discard(session.request_id)
         session._output_token_ids.clear()
         session._all_token_ids.clear()
-        # In-flight outputs from the previous segment were optimistically
-        # scheduled (async lookahead). Mark them stale so update_from_output
-        # drops them instead of underflowing num_output_placeholders
-        # (vLLM 0.27 a0c092ee72 removed async_tokens_to_discard). Seed in
-        # SCHEDULED-token units — num_in_flight_tokens matches what each
-        # pre-replacement frame will drain, so the counter reaches exactly
-        # zero; a placeholder-based seed swallowed valid new-segment frames
-        # whenever placeholder counts diverged from scheduled counts.
-        session.num_stale_output_tokens += int(getattr(session, "num_in_flight_tokens", 0) or 0)
-        session.num_output_placeholders = 0
-        session.spec_token_ids = []
         new_prompt = update.prompt_token_ids or ()
         session._all_token_ids.extend(new_prompt)
         session.num_computed_tokens = 0

vllm_omni/diffusion/compile.py

@@ -60,10 +60,7 @@ def regionally_compile(
     for name, submod in model.named_modules():
         if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
             # Compile the block compute while keeping nn.Module.__call__ hooks
-            # outside the compiled graph. NOTE: anything that wraps this
-            # callable must stay signature-transparent — cache-dit's
-            # BlockAdapter matches blocks by inspecting forward's parameter
-            # names and return annotation, which torch.compile preserves.
+            # outside the compiled graph.
             compiled_forwards.append(
                 (
                     submod,

vllm_omni/diffusion/data.py

@@ -27,7 +27,6 @@ from vllm_omni.diffusion.diffusion_kv.config import (
     DiffusionKVCacheMode,
     parse_diffusion_kv_cache_mode,
 )
-from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -683,9 +682,8 @@ class OmniDiffusionConfig:
     # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
 
     # LoRA parameters
-    lora_path: str | list[str] | None = None
-    lora_scale: float | list[float] = 1.0
-    lora_backend: LoRABackend = LoRABackend.PEFT  # available choices: ["peft", "distill"]
+    lora_path: str | None = None
+    lora_scale: float = 1.0
     max_cpu_loras: int | None = None
 
     output_type: str = "pil"

vllm_omni/diffusion/distributed/parallel_state.py

@@ -28,8 +28,6 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 
-import inspect
-
 import torch
 import torch.distributed
 import vllm.distributed.parallel_state as vllm_parallel_state
@@ -526,23 +524,13 @@ def init_vllm_model_parallel_group(
     local_rank: int,
     backend: str,
     group_name: str,
-    use_all2all: bool = False,
 ) -> vllm_parallel_state.GroupCoordinator:
-    # vLLM 0.27's MoE oracle (make_unquantized_moe_kernel ->
-    # get_ep_all2all_manager) asserts that the EP group's device communicator
-    # carries an all2all_manager, and vLLM only builds one when the group is
-    # created with use_all2all=True. Older vLLM has neither the kwarg nor the
-    # requirement, so forward it only when supported.
-    kwargs = {}
-    if use_all2all and "use_all2all" in inspect.signature(vllm_parallel_state.init_model_parallel_group).parameters:
-        kwargs["use_all2all"] = True
     return vllm_parallel_state.init_model_parallel_group(
         group_ranks=group_ranks,
         local_rank=local_rank,
         backend=backend,
         group_name=group_name,
         use_device_communicator=True,
-        **kwargs,
     )
 
 
@@ -966,7 +954,6 @@ def initialize_model_parallel(
             local_rank=get_world_group().local_rank,
             backend=backend,
             group_name="ep",
-            use_all2all=True,
         )
 
     init_dit_group(dit_parallel_size, backend)

vllm_omni/diffusion/layers/fused_moe.py

@@ -75,7 +75,7 @@ class FusedMoE:
         # Current vLLM FusedMoE handles output reduction internally.
         kwargs.pop("reduce_results", None)
         # FusedMoE is now a factory function; call it to get a MoERunner.
-        from vllm.model_executor.layers.fused_moe import FusedMoEFactory as _FusedMoE
+        from vllm.model_executor.layers.fused_moe import FusedMoE as _FusedMoE
 
         moe_runner = _FusedMoE(prefix=prefix, **kwargs)
 

vllm_omni/diffusion/lora/loader.py

@@ -1,461 +0,0 @@
-import os
-from collections import OrderedDict, defaultdict
-from collections.abc import Callable
-
-import torch
-from diffusers.loaders.lora_conversion_utils import (
-    _convert_non_diffusers_qwen_lora_to_diffusers,
-    _convert_non_diffusers_wan_lora_to_diffusers,
-)
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
-
-logger = init_logger(__name__)
-
-lora_convert_mapping: dict[str, Callable] = {
-    "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
-    "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
-    "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
-    "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
-    "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
-}
-
-
-def get_converter_by_pipeline(pipeline):
-    for pipeline_cls in pipeline.__class__.__mro__:
-        converter = lora_convert_mapping.get(pipeline_cls.__name__)
-        if converter is not None:
-            return converter
-    return None
-
-
-def _prepare_lora_delta(
-    lora_state_dict,
-    base_key: str,
-    param_to_weight_names: dict[str, list[str]] | None = None,
-    is_bias: bool = False,
-    lora_a_suffix: str = "lora_A.weight",
-    lora_b_suffix: str = "lora_B.weight",
-    lora_bias_suffix: str = "bias",
-):
-    used_keys = set()
-    # stacked_params_mapping                       param_to_weight_names
-    # [(".to_qkv", ".to_q.", "q")
-    # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
-    # (".to_qkv", ".to_v.", "v")]
-
-    # stacked_sd is to store packed parameter deltas (format: [str, list[torch.Tensor]])
-    # example: {".to_qkv": [delta_q, delta_k, delta_v]}
-    stacked_deltas = []
-    is_stacked_param = False
-    if param_to_weight_names is None:
-        param_to_weight_names = defaultdict(list)
-    for param_name, weight_names in param_to_weight_names.items():
-        if param_name not in base_key:
-            continue
-        is_stacked_param = True
-        # handle lora_a_key and lora_b_key together
-        if is_bias:
-            for weight_name in weight_names:
-                lora_bias_key = f"{base_key.replace(param_name, weight_name)}.{lora_bias_suffix}"
-                if lora_bias_key not in lora_state_dict:
-                    return None, used_keys
-                delta = lora_state_dict[lora_bias_key]
-                stacked_deltas.append(delta)
-                used_keys.add(lora_bias_key)
-        else:
-            for weight_name in weight_names:
-                lora_a_key = f"{base_key.replace(param_name, weight_name)}.{lora_a_suffix}"
-                lora_b_key = lora_a_key.replace(lora_a_suffix, lora_b_suffix)
-                if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
-                    return None, used_keys
-                a = lora_state_dict[lora_a_key]
-                b = lora_state_dict[lora_b_key]
-                delta = torch.matmul(b, a)
-                stacked_deltas.append(delta)
-                used_keys.add(lora_a_key)
-                used_keys.add(lora_b_key)
-        continue
-
-    if is_stacked_param:
-        return torch.concat(stacked_deltas), used_keys
-
-    if is_bias:
-        lora_bias_key = f"{base_key}.{lora_bias_suffix}"
-        if lora_bias_key not in lora_state_dict:
-            return None, used_keys
-        used_keys.add(lora_bias_key)
-        return lora_state_dict[lora_bias_key], used_keys
-
-    lora_a_key = f"{base_key}.{lora_a_suffix}"
-    lora_b_key = f"{base_key}.{lora_b_suffix}"
-    if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
-        return None, used_keys
-    a = lora_state_dict[lora_a_key]
-    b = lora_state_dict[lora_b_key]
-    used_keys.add(lora_a_key)
-    used_keys.add(lora_b_key)
-    return torch.matmul(b, a), used_keys
-
-
-def _load_lora_state_dict(
-    pretrained_model_name_or_path: str,
-    weights_name: str | None = None,
-    use_safetensors: bool = True,
-    subfolder: str | None = None,
-    cache_dir: str | None = None,
-    force_download: bool = False,
-    local_files_only: bool | None = None,
-):
-    # first we try to load it from local disk
-    if use_safetensors and pretrained_model_name_or_path.endswith(".safetensors"):
-        return load_file(pretrained_model_name_or_path)
-
-    if os.path.isdir(pretrained_model_name_or_path):
-        model_file = pretrained_model_name_or_path
-        if subfolder:
-            model_file = os.path.join(model_file, subfolder)
-        if weights_name is None:
-            raise ValueError("weights_name is required when loading from a directory")
-        model_file = os.path.join(model_file, weights_name)
-        return load_file(model_file)
-
-    # finally, we try to load it from the internet
-    try:
-        model_file = hf_hub_download(
-            pretrained_model_name_or_path,
-            filename=weights_name,
-            subfolder=subfolder,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            local_files_only=local_files_only,
-        )
-        return load_file(model_file)
-    except Exception as e:
-        logger.error(f"Failed to download {weights_name} from {pretrained_model_name_or_path}: {e}")
-        raise e
-
-
-def _remap_state_dict_keys(sd, rules):
-    """
-    Remap keys in a state_dict by sequentially applying a list of substring substitution rules.
-
-    This utility function transforms parameter names by applying a sequence of
-    (old_substring, new_substring) pairs. For each key in the input dictionary,
-    **all rules are evaluated in the given order**, and every rule whose
-    `old_substring` is found within the current key triggers a replacement.
-    Substitutions are cumulative: the key is updated after each match, and
-    subsequent rules operate on the result of previous replacements.
-
-    The order of rules matters; users are responsible for arranging them to
-    achieve the desired final key. Keys that do not match any rule are copied
-    unchanged.
-
-    Args:
-        sd (dict[str, Any]): Original state_dict mapping parameter names to tensors.
-        rules (List[Tuple[str, str]]): A list of (old_substring, new_substring) pairs.
-            Rules are applied sequentially, and **all matching rules are executed**.
-
-    Returns:
-        dict[str, Any]: A new state_dict with remapped keys. Values are unchanged.
-    """
-    new_sd = OrderedDict()
-    for k, v in sd.items():
-        new_key = k
-        matched = False
-        for p1, p2 in rules:
-            if p1 not in k:
-                continue
-            new_key = new_key.replace(p1, p2)
-            matched = True
-        if matched:
-            new_sd[new_key] = v
-        else:
-            new_sd[k] = v
-    return new_sd
-
-
-def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-    """Fold Diffusers-format LoRA alpha values into their B matrices."""
-    scaled_state_dict = dict(state_dict)
-    for alpha_key in (key for key in state_dict if key.endswith(".alpha")):
-        base_key = alpha_key[: -len(".alpha")]
-        lora_a_key = f"{base_key}.lora_A.weight"
-        lora_b_key = f"{base_key}.lora_B.weight"
-        if lora_a_key not in state_dict or lora_b_key not in state_dict:
-            raise ValueError(f"LoRA alpha key {alpha_key!r} does not have matching lora_A and lora_B weights.")
-
-        rank = state_dict[lora_a_key].shape[0]
-        alpha = state_dict[alpha_key]
-
-        scaled_state_dict[lora_b_key] = state_dict[lora_b_key] * (alpha.item() / rank)
-        scaled_state_dict.pop(alpha_key)
-
-    return scaled_state_dict
-
-
-class LoraLoaderMixin:
-    transformer_name = "transformer"
-
-    # Lazy initialization to avoid MRO issues: __init__ may not be called
-    # when mixin with nn.Module
-    @property
-    def lora_loaded(self):
-        if not hasattr(self, "_lora_loaded"):
-            self._lora_loaded = {}
-        return self._lora_loaded
-
-    @lora_loaded.setter
-    def lora_loaded(self, value):
-        self._lora_loaded = value
-
-    @classmethod
-    def load_lora_into_module(
-        cls,
-        state_dict,
-        module,
-        prefix: str = "transformer",
-        lora_a_suffix: str = "lora_A.weight",
-        lora_b_suffix: str = "lora_B.weight",
-        lora_bias_suffix: str = "bias",
-    ):
-        lora_loaded_keys = set()
-
-        # prepare stacked parameters mapping for later use
-        # stacked_params_mapping                       param_to_weight_names
-        # [(".to_qkv", ".to_q.", "q")
-        # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
-        # (".to_qkv", ".to_v.", "v")]
-        param_to_weight_names = defaultdict(list)
-        if hasattr(module, "stacked_params_mapping"):
-            for param_name, weight_name, _ in module.stacked_params_mapping:
-                param_to_weight_names[param_name].append(weight_name)
-
-        for name, params in module.named_parameters(prefix):
-            is_bias = False
-            if name.endswith(".bias"):
-                base_key = name[: -len(".bias")]
-                is_bias = True
-            elif name.endswith(".weight"):
-                base_key = name[: -len(".weight")]
-            else:
-                continue
-
-            delta, used_keys = _prepare_lora_delta(
-                state_dict,
-                base_key,
-                param_to_weight_names,
-                is_bias,
-                lora_a_suffix,
-                lora_b_suffix,
-                lora_bias_suffix,
-            )
-            if delta is None:
-                continue
-            lora_loaded_keys.update(used_keys)
-
-            delta = delta.to(device=params.device, dtype=params.dtype)
-            with torch.no_grad():
-                params.add_(delta)
-            del delta
-
-        missing_keys = set(state_dict.keys()) - lora_loaded_keys
-        for k in missing_keys:
-            logger.warning(f"Missing loading lora key: {k}")
-
-        logger.info(f"{len(lora_loaded_keys)} lora keys loaded into {module.__class__.__name__}.")
-        return lora_loaded_keys
-
-    @classmethod
-    def unload_module_lora(
-        cls,
-        state_dict,
-        module,
-        prefix: str = "transformer",
-        lora_a_suffix: str = "lora_A.weight",
-        lora_b_suffix: str = "lora_B.weight",
-        lora_bias_suffix: str = "bias",
-    ):
-        lora_unloaded_keys = set()
-
-        param_to_weight_names = defaultdict(list)
-        if hasattr(module, "stacked_params_mapping"):
-            for param_name, weight_name, _ in module.stacked_params_mapping:
-                param_to_weight_names[param_name].append(weight_name)
-
-        for name, param in module.named_parameters(prefix):
-            is_bias = False
-            if name.endswith(".bias"):
-                base_key = name[: -len(".bias")]
-                is_bias = True
-            elif name.endswith(".weight"):
-                base_key = name[: -len(".weight")]
-            else:
-                continue
-
-            delta, used_keys = _prepare_lora_delta(
-                state_dict,
-                base_key,
-                param_to_weight_names,
-                is_bias,
-                lora_a_suffix,
-                lora_b_suffix,
-                lora_bias_suffix,
-            )
-            if delta is None:
-                continue
-            lora_unloaded_keys.update(used_keys)
-
-            delta = delta.to(device=param.device, dtype=param.dtype)
-            with torch.no_grad():
-                param.sub_(delta)
-            del delta
-
-        missing_keys = set(state_dict.keys()) - lora_unloaded_keys
-        for k in missing_keys:
-            logger.warning(f"Missing unloading lora key: {k}")
-
-        logger.info(f"{len(lora_unloaded_keys)} lora keys unloaded from {module.__class__.__name__}.")
-        return lora_unloaded_keys
-
-
-class QwenImageLoraLoaderMixin(LoraLoaderMixin):
-    def load_lora_weights(
-        self,
-        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
-        adapter_name: str | None = None,
-    ):
-        if adapter_name in self.lora_loaded:
-            return
-
-        if isinstance(pretrained_model_name_or_path_or_dict, dict):
-            state_dict = pretrained_model_name_or_path_or_dict
-        else:
-            state_dict = _load_lora_state_dict(pretrained_model_name_or_path_or_dict)
-
-        has_alpha = any(k.endswith(".alpha") for k in state_dict)
-        is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
-        is_diffusers_lora_format = any(k.endswith(".lora_A.weight") for k in state_dict)
-        if has_alpha and is_diffusers_lora_format:
-            state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
-
-        if has_alpha or is_non_diffusers_format:
-            converter = get_converter_by_pipeline(self)
-            if converter is None:
-                raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
-
-            state_dict = converter(state_dict)
-
-        state_dict = _remap_state_dict_keys(state_dict, [(".to_out.0.", ".to_out.")])
-
-        self.load_lora_into_module(
-            state_dict,
-            self.transformer,
-            prefix=self.transformer_name,
-        )
-
-        self.lora_loaded[adapter_name] = state_dict
-
-    def unload_lora_weights(self, adapter_name: str):
-        if adapter_name not in self.lora_loaded:
-            return
-
-        state_dict = self.lora_loaded[adapter_name]
-
-        transformer = get_transformer_from_pipeline(self)
-        self.unload_module_lora(state_dict, transformer, prefix=self.transformer_name)
-
-        del self.lora_loaded[adapter_name]
-
-
-class WanLoraLoaderMixin(LoraLoaderMixin):
-    transformer_2_name = "transformer_2"
-
-    def load_lora_weights(
-        self,
-        pretrained_model_name_or_path: str | list[str],
-        adapter_name: str | None = None,
-    ):
-        """Load LoRA weights into the pipeline's transformer(s).
-
-        Args:
-            pretrained_model_name_or_path: Path(s) to concrete .safetensors LoRA files.
-                - str: a single .safetensors file. Only valid for Wan2.1
-                  (single-transformer pipelines).
-                - list[str]: one file per transformer, matched **by position** to
-                  [transformer, transformer_2]. Required for Wan2.2 MoE.
-            adapter_name: Name to register the adapter under.
-        """
-        if adapter_name in self.lora_loaded:
-            return
-
-        # Normalize to list[(lora_path, target_module_name)].
-        lora_paths = None
-        if isinstance(pretrained_model_name_or_path, str):
-            lora_paths = [pretrained_model_name_or_path]
-        else:
-            lora_paths = list(pretrained_model_name_or_path)
-
-        target_modules = (
-            [self.transformer_name, self.transformer_2_name] if self.has_transformer_2 else [self.transformer_name]
-        )
-        if len(lora_paths) != len(target_modules):
-            raise ValueError(
-                f"{self.__class__.__name__} expects {len(target_modules)} LoRA file(s) "
-                f"(target_modules={target_modules}, paths={lora_paths}), got {len(lora_paths)}. Pass a list of concrete "  # noqa
-                f".safetensors files, one per transformer."
-            )
-
-        module_to_lora_sd = dict()
-        for lora_path, target_module_name in zip(lora_paths, target_modules):
-            state_dict = _load_lora_state_dict(lora_path)
-            is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
-            if is_non_diffusers_format:
-                converter = get_converter_by_pipeline(self)
-                if converter is None:
-                    raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
-
-                state_dict = converter(state_dict)
-
-            state_dict = _remap_state_dict_keys(
-                state_dict,
-                [
-                    (".ffn.net.0.", ".ffn.net_0."),
-                    (".ffn.net.2.", ".ffn.net_2."),
-                    (
-                        ".to_out.0.",
-                        ".to_out.",
-                    ),
-                    # '.bias_B.bias' is a horrible name in diffuser's.
-                    # So we remap it to '.bias'
-                    (".lora_B.bias", ".bias"),
-                ],
-            )
-
-            module = get_transformer_from_pipeline(self, transformer_name=target_module_name)
-            if module is None:
-                logger.warning(
-                    f"Skip LoRA {lora_path}: target module '{target_module_name}' "
-                    f"is not loaded in pipeline {self.__class__.__name__} "
-                    f"(e.g. disabled by boundary_ratio)."
-                )
-                continue
-
-            self.load_lora_into_module(state_dict, module, prefix=self.transformer_name, lora_bias_suffix="bias")
-            module_to_lora_sd[target_module_name] = state_dict
-
-        self.lora_loaded[adapter_name] = module_to_lora_sd
-
-    def unload_lora_weights(self, adapter_name: str):
-        if adapter_name not in self.lora_loaded:
-            return
-
-        module_to_lora_sd = self.lora_loaded[adapter_name]
-        for module_name, lora_sd in module_to_lora_sd.items():
-            module = get_transformer_from_pipeline(self, transformer_name=module_name)
-            self.unload_module_lora(lora_sd, module, prefix=self.transformer_name)
-
-        del self.lora_loaded[adapter_name]

vllm_omni/diffusion/lora/manager.py

@@ -4,7 +4,6 @@
 import math
 import time
 from collections import OrderedDict
-from enum import Enum
 from typing import get_args
 
 import torch
@@ -34,11 +33,6 @@ from vllm_omni.lora.utils import stable_lora_int_id
 logger = init_logger(__name__)
 
 
-class LoRABackend(str, Enum):
-    PEFT = "peft"
-    DISTILL = "distill"
-
-
 class DiffusionLoRAManager:
     """Manager for LoRA adapters in diffusion models.
 

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -36,22 +36,6 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
-# Preference chain for CUDA flash attention in diffusers, best first. Shared
-# with the accuracy tests' diffusers reference runs so both sides of a
-# similarity comparison resolve to the same backend on any given image
-# (hub kernels may lack a build variant for the image's torch).
-CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS: list[str] = [
-    "_flash_3_hub",
-    "_flash_3_varlen_hub",
-    "_flash_3",
-    "_flash_varlen_3",
-    "flash_hub",
-    "flash_varlen_hub",
-    "flash",
-    "flash_varlen",
-    "_native_flash",
-]
-
 _DIFFUSERS_CONFIG_LOAD_KWARGS = {
     "cache_dir",
     "dduf_entries",
@@ -295,7 +279,19 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                     )
                     attention_backend_attempts.append("native")
                 else:
-                    attention_backend_attempts.extend(CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS)
+                    attention_backend_attempts.extend(
+                        [
+                            "_flash_3_hub",
+                            "_flash_3_varlen_hub",
+                            "_flash_3",
+                            "_flash_varlen_3",
+                            "flash_hub",
+                            "flash_varlen_hub",
+                            "flash",
+                            "flash_varlen",
+                            "_native_flash",
+                        ]
+                    )
             case "SAGE_ATTN":
                 attention_backend_attempts.extend(["sage_hub", "sage", "sage", "sage_varlen"])
             case "ASCEND":

vllm_omni/diffusion/models/lingbot_video/request_utils.py

@@ -299,21 +299,16 @@ def resolve_lingbot_output_dimensions(
     sampling_width: Any = None,
     sampling_height: Any = None,
     prompt_fields: Mapping[str, Any] | None = None,
-    extra_fields: Mapping[str, Any] | None = None,
     default_width: int = 480,
     default_height: int = 480,
 ) -> tuple[int, int]:
     """Resolve the ``(width, height)`` that LingBot will use for generation."""
     prompt_fields = prompt_fields or {}
-    extra_fields = extra_fields or {}
-    # Online requests deliver `extra_params` through sampling.extra_args; the
-    # offline path embeds the same fields in the prompt dict. Consult both,
-    # extra_args first — the same precedence fps and num_frames use.
-    requested_width = _first_not_none(_pick(extra_fields, "width"), _pick(prompt_fields, "width"))
-    requested_height = _first_not_none(_pick(extra_fields, "height"), _pick(prompt_fields, "height"))
-    size = _first_not_none(_pick(extra_fields, "size"), _pick(prompt_fields, "size"))
-    resolution = _first_not_none(_pick(extra_fields, "resolution"), _pick(prompt_fields, "resolution"))
-    ratio = _first_not_none(_pick(extra_fields, "ratio"), _pick(prompt_fields, "ratio"))
+    requested_width = _pick(prompt_fields, "width")
+    requested_height = _pick(prompt_fields, "height")
+    size = _pick(prompt_fields, "size")
+    resolution = _pick(prompt_fields, "resolution")
+    ratio = _pick(prompt_fields, "ratio")
     if requested_width is not None or requested_height is not None:
         width, height = requested_width, requested_height
     elif size is not None or resolution is not None or ratio is not None:
@@ -419,7 +414,6 @@ def normalize_lingbot_request(
         sampling_width=sampling.width,
         sampling_height=sampling.height,
         prompt_fields=prompt_fields,
-        extra_fields=extra_args,
         default_width=default_width,
         default_height=default_height,
     )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -25,7 +25,6 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.lora.loader import QwenImageLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -265,11 +264,7 @@ def apply_rotary_emb_qwen(
 
 
 class QwenImagePipeline(
-    nn.Module,
-    QwenImageCFGParallelMixin,
-    DiffusionPipelineProfilerMixin,
-    SupportsComponentDiscovery,
-    QwenImageLoraLoaderMixin,
+    nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
 ):
     supports_request_batch = True
     _dit_modules: ClassVar[list[str]] = ["transformer"]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -25,7 +25,6 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
-from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -286,7 +285,6 @@ class Wan22Pipeline(
     DenoiseProgressMixin,
     DiffusionPipelineProfilerMixin,
     SupportsComponentDiscovery,
-    WanLoraLoaderMixin,
 ):
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -25,7 +25,6 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
-from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -148,7 +147,6 @@ class Wan22I2VPipeline(
     DiffusionPipelineProfilerMixin,
     DenoiseProgressMixin,
     SupportsComponentDiscovery,
-    WanLoraLoaderMixin,
 ):
     """
     Wan2.2 Image-to-Video Pipeline.

vllm_omni/diffusion/utils/tf_utils.py

@@ -70,10 +70,9 @@ def find_module_with_attr(model, attr_name="transformer"):
     return None
 
 
-def get_transformer_from_pipeline(pipeline: Any, transformer_name: str = "transformer"):
-    """Return a named transformer from a pipeline or pipeline wrapper."""
-    pipe = find_module_with_attr(pipeline, attr_name=transformer_name)
+def get_transformer_from_pipeline(pipeline: Any):
+    pipe = find_module_with_attr(pipeline, attr_name="transformer")
 
     if pipe is not None:
-        return getattr(pipe, transformer_name)
+        return pipe.transformer
     return None

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -49,7 +49,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
-from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -388,31 +388,14 @@ class DiffusionWorker:
         """Initialize the LoRA manager for this worker."""
         if self.model_runner.pipeline is None:
             return
-
-        lora_path = self.od_config.lora_path
-        if isinstance(lora_path, list) and len(lora_path) == 1:
-            lora_path = lora_path[0]
-
-        lora_backend = self.od_config.lora_backend
-        if lora_backend == LoRABackend.PEFT:
-            self.lora_manager = DiffusionLoRAManager(
-                pipeline=self.model_runner.pipeline,
-                device=self.device,
-                dtype=self.od_config.dtype,
-                max_cached_adapters=self.od_config.max_cpu_loras,
-                lora_path=lora_path,
-                lora_scale=self.od_config.lora_scale,
-            )
-        elif lora_backend == LoRABackend.DISTILL:
-            pipeline = self.model_runner.pipeline
-            if hasattr(pipeline, "load_lora_weights"):
-                if self.od_config.lora_scale > 1.0:
-                    logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.")
-                pipeline.load_lora_weights(lora_path)
-            else:
-                logger.warning("Pipeline does not support loading distilled LoRA weights for now.")
-        else:
-            raise ValueError(f"Unknown LoRA backend: {lora_backend}. Available choices: {LoRABackend.__members__}")
+        self.lora_manager = DiffusionLoRAManager(
+            pipeline=self.model_runner.pipeline,
+            device=self.device,
+            dtype=self.od_config.dtype,
+            max_cached_adapters=self.od_config.max_cpu_loras,
+            lora_path=self.od_config.lora_path,
+            lora_scale=self.od_config.lora_scale,
+        )
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop profiling for this GPU worker.

vllm_omni/engine/arg_utils.py

@@ -104,16 +104,10 @@ def register_omni_models_to_vllm():
 
     _register_omni_hf_configs()
 
-    # Unconditionally (re)register every omni arch into the upstream global
-    # ModelRegistry: vLLM 0.27 added some omni archs (e.g.
-    # Qwen3OmniMoeForConditionalGeneration) to its own registry, but resolves
-    # them to upstream classes that lack omni features (e.g.
-    # buffer_realtime_audio for the realtime endpoint). Registering here
-    # overwrites those entries with the vllm-omni classes, which is the
-    # behavior entrypoints resolving through the global registry rely on
-    # (and matches OmniModelRegistry's omni-wins ordering).
+    supported_archs = ModelRegistry.get_supported_archs()
     for arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items():
-        ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
+        if arch not in supported_archs:
+            ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
 
     # Register omni-specific reasoning parsers (e.g., step_audio).
     import vllm_omni.reasoning  # noqa: F401
@@ -477,9 +471,6 @@ class OrchestratorArgs:
     num_gpus: int | None = None
     model_class_name: str | None = None
     diffusion_load_format: str | None = None
-    lora_path: list[str] | None = None
-    lora_backend: str | None = None
-    lora_scale: float | None = None
     diffusers_load_kwargs: str = "{}"
     diffusers_call_kwargs: str = "{}"
     ulysses_degree: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -539,7 +539,6 @@ class AsyncOmniEngine:
         if not isinstance(mm_data, dict) or not mm_data:
             return
 
-        from vllm.config.multimodal import _get_mm_hasher_algorithm
         from vllm.multimodal.hasher import MultiModalHasher
 
         existing_uuids = prompt.get("multi_modal_uuids")
@@ -566,7 +565,6 @@ class AsyncOmniEngine:
                     base_uuid = None
                 else:
                     base_uuid = MultiModalHasher.hash_kwargs(
-                        _get_mm_hasher_algorithm(),
                         model_id=model_id,
                         **{modality: item},
                     )
@@ -1050,9 +1048,6 @@ class AsyncOmniEngine:
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),
             "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
-            "lora_path": kwargs.get("lora_path", None),
-            "lora_scale": kwargs.get("lora_scale", 1.0),
-            "lora_backend": kwargs.get("lora_backend", "peft"),
             "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
             "worker_extension_cls": kwargs.get("worker_extension_cls", None),
             "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]),
@@ -1256,9 +1251,6 @@ class AsyncOmniEngine:
                 if lora_scale is not None:
                     if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
                         cfg.engine_args.lora_scale = lora_scale
-                if kwargs.get("lora_backend") is not None:
-                    if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None:
-                        cfg.engine_args.lora_backend = kwargs["lora_backend"]
                 if (
                     kwargs.get("diffusion_attention_config") is not None
                     or kwargs.get("diffusion_attention_backend") is not None

vllm_omni/engine/stage_engine_startup.py

@@ -896,59 +896,23 @@ def connect_remote_diffusion_proc(
 def scoped_spawn_device_env(
     stage_visible_devices: str | None,
     spawn_device_lock: threading.Lock | None,
-    *,
-    stage_id: int | None = None,
-    replica_id: int | None = None,
 ) -> Iterator[None]:
     """Briefly scope device visibility while spawning a stage subprocess."""
-    with _scoped_replica_compile_cache_env(stage_id, replica_id):
-        if stage_visible_devices is None or spawn_device_lock is None:
-            yield
-            return
-
-        device_control_env = current_omni_platform.device_control_env_var
-        with spawn_device_lock:
-            previous_visible_devices = os.environ.get(device_control_env)
-            try:
-                current_omni_platform.set_device_control_env_var(stage_visible_devices)
-                yield
-            finally:
-                if previous_visible_devices is None:
-                    current_omni_platform.unset_device_control_env_var()
-                else:
-                    current_omni_platform.set_device_control_env_var(previous_visible_devices)
-
-
-_COMPILE_CACHE_ENV_LOCK = threading.Lock()
-
-
-@contextlib.contextmanager
-def _scoped_replica_compile_cache_env(stage_id: int | None, replica_id: int | None) -> Iterator[None]:
-    """Give additional replicas an isolated VLLM_CACHE_ROOT at spawn time.
-
-    Replicas of the same stage compile the same model concurrently; vLLM
-    0.27's AOT compile (TritonBundler) races on the shared cache dir and a
-    replica can then die with 'Cubin file saved by TritonBundler not found'
-    (build 2954, Multi-Replica perf). Replica 0 keeps the default shared
-    cache so single-replica deployments stay warm; every further replica
-    compiles into its own subdirectory.
-    """
-    if not replica_id or replica_id <= 0:
+    if stage_visible_devices is None or spawn_device_lock is None:
         yield
         return
 
-    base = os.environ.get("VLLM_CACHE_ROOT") or os.path.expanduser("~/.cache/vllm")
-    isolated = os.path.join(base, f"omni_stage{stage_id if stage_id is not None else 'x'}_replica{replica_id}")
-    with _COMPILE_CACHE_ENV_LOCK:
-        previous = os.environ.get("VLLM_CACHE_ROOT")
-        os.environ["VLLM_CACHE_ROOT"] = isolated
+    device_control_env = current_omni_platform.device_control_env_var
+    with spawn_device_lock:
+        previous_visible_devices = os.environ.get(device_control_env)
         try:
+            current_omni_platform.set_device_control_env_var(stage_visible_devices)
             yield
         finally:
-            if previous is None:
-                os.environ.pop("VLLM_CACHE_ROOT", None)
+            if previous_visible_devices is None:
+                current_omni_platform.unset_device_control_env_var()
             else:
-                os.environ["VLLM_CACHE_ROOT"] = previous
+                current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
 
 @contextlib.contextmanager
@@ -1043,9 +1007,7 @@ def _launch_omni_core_engines(
                 # an OmniCoordClientForStage and heartbeats to the coordinator.
                 from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
 
-                with scoped_spawn_device_env(
-                    stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id
-                ):
+                with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
                     local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager(
                         local_engine_count=local_engine_count,
                         start_index=start_index,
@@ -1060,9 +1022,7 @@ def _launch_omni_core_engines(
                         omni_replica_base_id=replica_id,
                     )
             else:
-                with scoped_spawn_device_env(
-                    stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id
-                ):
+                with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
                     local_engine_manager = CoreEngineProcManager(
                         local_engine_count=local_engine_count,
                         start_index=start_index,
@@ -1144,7 +1104,7 @@ def launch_stage_replica(
     addresses = get_engine_zmq_addresses(vllm_config)
     handshake_address = get_open_zmq_ipc_path()
     engines_to_handshake = [CoreEngine(index=0, local=True)]
-    with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id):
+    with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
         engine_manager = StageEngineCoreProcManager(
             local_engine_count=1,
             start_index=0,

vllm_omni/entrypoints/cli/serve.py

@@ -164,13 +164,6 @@ class OmniServeCommand(CLISubcommand):
                     "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
                 )
 
-        lora_backend = getattr(args, "lora_backend", None)
-        lora_path = getattr(args, "lora_path", None)
-        if lora_backend == "distill" and not lora_path:
-            raise ValueError("--lora-backend distill requires --lora-path")
-        if (lora_backend or "peft") == "peft" and isinstance(lora_path, list) and len(lora_path) > 1:
-            raise ValueError("--lora-backend peft accepts only one startup LoRA path")
-
         # --omni-lb-policy is validated against the LoadBalancingPolicy enum.
         omni_lb_policy = getattr(args, "omni_lb_policy", None)
         if omni_lb_policy is not None:
@@ -434,32 +427,6 @@ class OmniServeCommand(CLISubcommand):
                 "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
             ),
         )
-        omni_config_group.add_argument(
-            "--lora-path",
-            type=str,
-            nargs="+",
-            default=None,
-            help=(
-                "LoRA checkpoint path(s) loaded when the diffusion server starts. "
-                "The distill backend accepts one file per pipeline transformer."
-            ),
-        )
-        omni_config_group.add_argument(
-            "--lora-backend",
-            type=str,
-            choices=["peft", "distill"],
-            default=None,
-            help=(
-                "Diffusion LoRA loading backend. 'distill' fuses checkpoint files "
-                "into the base model at server startup; 'peft' uses the adapter manager."
-            ),
-        )
-        omni_config_group.add_argument(
-            "--lora-scale",
-            type=float,
-            default=None,
-            help="Scale for a startup PEFT LoRA. Distilled LoRAs are fused at their checkpoint scale.",
-        )
         omni_config_group.add_argument(
             "--diffusion-compile-granularity",
             choices=["regional", "full"],

vllm_omni/entrypoints/openai/api_server.py

@@ -152,7 +152,6 @@ from vllm_omni.entrypoints.openai.video_api_utils import (
     decode_input_reference,
 )
 from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI
-from vllm_omni.entrypoints.utils import PureDiffusionLauncherAdapter
 from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.utils.forced_aligner import build_forced_aligner_config
@@ -550,16 +549,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
                 "Starting vLLM API server (pure diffusion mode) on %s",
                 listen_address,
             )
-            # The vLLM 0.27 launcher's shutdown path reads
-            # engine_client.vllm_config.shutdown_timeout; AsyncOmni.vllm_config
-            # is None for pure diffusion (no comprehension stage), which would
-            # crash handle_shutdown and hang teardown. Wrap app.state.engine_client
-            # with a shim that only overrides vllm_config and forwards everything
-            # else (get_vllm_config still returns None for pure-diffusion detection).
-            app.state.engine_client = PureDiffusionLauncherAdapter(
-                engine_client,
-                shutdown_timeout=getattr(args, "shutdown_timeout", 0),
-            )
         else:
             logger.info(
                 "Starting vLLM API server %d on %s",
@@ -996,9 +985,8 @@ async def omni_init_app_state(
     )
 
     # Warm up chat template processing to avoid first-request latency
-    # Upstream f5ffc59b6a removed OpenAIServingChat.warmup() and moved the
-    # warmup onto the renderer (OnlineRenderer.warmup()); mirror upstream.
-    state.online_renderer.warmup()
+    if state.openai_serving_chat is not None:
+        state.openai_serving_chat.warmup()
 
     state.openai_serving_completion = (
         OpenAIServingCompletion(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -83,7 +83,6 @@ from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.multimodal.media.connector import MediaConnector
 from vllm.outputs import RequestOutput
-from vllm.parser.mistral import MistralToolCall
 from vllm.reasoning import ReasoningParser
 from vllm.renderers import BaseRenderer, merge_kwargs
 from vllm.renderers.inputs import TokPrompt
@@ -97,6 +96,7 @@ from vllm.tokenizers.mistral import (
     validate_request_params,
 )
 from vllm.tool_parsers import ToolParser
+from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.tool_parsers.streaming import extract_required_tool_call_streaming
 from vllm.utils.collection_utils import as_list
 from vllm.v1.engine.exceptions import EngineDeadError

vllm_omni/entrypoints/utils.py

@@ -905,31 +905,3 @@ def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool):
         params.skip_clone = True
     params.output_kind = target_type
     return params
-
-
-class PureDiffusionLauncherAdapter:
-    """vLLM launcher compatibility shim for pure-diffusion mode.
-
-    The upstream launcher's shutdown path reads
-    ``app.state.engine_client.vllm_config.shutdown_timeout``
-    (vllm/entrypoints/launcher.py), but ``AsyncOmni.vllm_config`` returns
-    ``None`` when the pipeline has no comprehension stage (pure diffusion),
-    which crashes ``handle_shutdown`` with AttributeError and hangs server
-    teardown (workers force-killed, spurious resource_tracker noise).
-
-    This adapter only overrides the ``vllm_config`` property with a minimal
-    fallback carrying ``shutdown_timeout`` and forwards every other attribute
-    to the wrapped engine client, so the pure-diffusion detection
-    (``get_vllm_config()`` still returns ``None``) is unaffected.
-    """
-
-    def __init__(self, engine_client: Any, shutdown_timeout: float) -> None:
-        object.__setattr__(self, "_wrapped", engine_client)
-        object.__setattr__(self, "_shutdown_timeout", float(shutdown_timeout))
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._wrapped, name)
-
-    @property
-    def vllm_config(self) -> Any:
-        return types.SimpleNamespace(shutdown_timeout=self._shutdown_timeout)

vllm_omni/model_executor/models/dynin_omni/dynin_omni.py

@@ -431,10 +431,7 @@ class DyninOmniMultiModalProcessor(BaseMultiModalProcessor[DyninOmniProcessingIn
         mm_items = inputs.mm_data_items
 
         with timing_ctx.record("get_mm_hashes"):
-            mm_hashes = inputs.get_mm_hashes(
-                self.info.model_id,
-                self.info.ctx.get_mm_config().mm_hasher_algorithm,
-            )
+            mm_hashes = inputs.get_mm_hashes(self.info.model_id)
 
         tokenizer = self.info.ctx.tokenizer
         prompt_token_ids = self._encode_prompt_to_token_ids(prompt, tokenizer)

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -29,7 +29,7 @@ from vllm.distributed import (
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory as SharedFusedMoE
+from vllm.model_executor.layers.fused_moe import FusedMoE as SharedFusedMoE
 from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -49,12 +49,13 @@ from vllm.model_executor.models.utils import (
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.config import (
+    is_interleaved,
     patch_rope_parameters,
     set_default_rope_theta,
 )
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, is_interleaved
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
 

vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py

@@ -33,7 +33,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
-    FusedMoEFactory,
+    FusedMoE,
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -332,7 +332,7 @@ class BailingMoeV2SparseMoeBlock(nn.Module):
         else:
             self.shared_experts = None
 
-        self.experts = FusedMoEFactory(
+        self.experts = FusedMoE(
             shared_experts=self.shared_experts,
             num_experts=config.num_experts,
             top_k=config.num_experts_per_tok,

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -76,14 +76,7 @@ class Qwen2_5OmniForConditionalGeneration(
         talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
         self.talker_config = talker_config
 
-        # See qwen3_omni for the full rationale. model_stage is only set by omni's
-        # staged startup; this class is now also constructed for non-staged runs
-        # (its arch overrides upstream's in the registry), where reading the
-        # attribute directly raises AttributeError during engine-core init.
-        # Default to the thinker — the text-generation stage. Note there is no
-        # else-branch below, so an unset stage would otherwise build no submodule
-        # at all and fail later at self.thinker.
-        self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "thinker"
+        self.model_stage = vllm_config.model_config.model_stage
         if self.model_stage == "thinker":
             # Initialize thinker model (multimodal processing)
             self.thinker = init_vllm_registered_model(

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -296,23 +296,9 @@ class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
     def get_processor_data(self) -> Mapping[str, object]:
         from transformers.video_utils import VideoMetadata
 
-        # Read through self.get(), not self.data. vLLM now wraps loaded media in
-        # MediaWithBytes (multimodal/media/base.py) to retain the original bytes
-        # for hashing, and VideoProcessorItems._unwrap is what strips it back to
-        # the raw frames. Iterating self.data skips that unwrap, so the HF
-        # processor received [MediaWithBytes, ...] where it expects frame arrays.
-        # transformers' make_batched_videos recognises the wrapper as neither a
-        # video nor a sequence, silently builds an empty list, and then indexes
-        # it -- surfacing as a 400 on every video request:
-        #     File "transformers/video_utils.py", in convert_pil_frames_to_video
-        #       if not (isinstance(videos[0], (list, tuple)) and ...
-        #     IndexError: list index out of range
-        # self.get() is also correct on the pre-wrapper vLLM, where it is a
-        # plain self.data[index].
-        items = [self.get(index) for index in range(self.get_count())]
-        videos = [frames for frames, _ in items]
+        videos = [frames for frames, _ in self.data]
         video_metadata = [
-            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in items
+            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
         ]
         return {"videos": videos, "video_metadata": video_metadata}
 
@@ -636,10 +622,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         )
 
         with timing_ctx.record("get_mm_hashes"):
-            mm_hashes = inputs.get_mm_hashes(
-                self.info.model_id,
-                self.info.ctx.get_mm_config().mm_hasher_algorithm,
-            )
+            mm_hashes = inputs.get_mm_hashes(self.info.model_id)
 
         mm_cache_keys = self._paired_cache_keys(mm_hashes, aiv_pairs)
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -127,25 +127,8 @@ class Qwen3OmniMoeForConditionalGeneration(
         code2wav_config: Qwen3OmniMoeCode2WavConfig = config.code2wav_config
         self.code2wav_config = code2wav_config
 
-        # Determine model stage. model_stage is injected by omni's staged
-        # startup (engine/stage_init_utils.py); a plain, non-staged vLLM run
-        # never sets it, and reading it directly raised
-        #   AttributeError: 'ModelConfig' object has no attribute 'model_stage'
-        # killing the engine core before it became ready. That path became
-        # reachable once omni archs began overriding upstream's in the global
-        # registry, so this class is now also constructed for non-staged runs
-        # such as the vLLM-text perf benchmark.
-        #
-        # Default to the thinker: it is the text-generation stage, which is what
-        # a non-staged run of this model is asking for. Same shape as
-        # dynin_omni ("token2text") and glm_tts ("glm_tts"). An explicitly wrong
-        # value still reaches the ValueError below rather than being silently
-        # accepted.
-        self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "thinker"
-        # Staged startup always injects model_stage; its absence means a plain
-        # vLLM run with no talker stage downstream, so no one consumes captured
-        # thinker layers and the forward must return what stock vLLM expects.
-        self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None
+        # Determine model stage
+        self.model_stage = vllm_config.model_config.model_stage
 
         if self.model_stage == "thinker":
             self.use_async_omni_output = True
@@ -403,21 +386,17 @@ class Qwen3OmniMoeForConditionalGeneration(
                 inputs_embeds = inputs_embeds.to(thinker_dev)
 
             # Run thinker forward
-            # If talker expects a specific intermediate layer, capture it here.
-            # Only staged runs have a talker stage to consume the capture; in a
-            # plain vLLM run capturing would waste a clone per layer per step
-            # and make the thinker return a tuple stock vLLM cannot handle.
+            # If talker expects a specific intermediate layer, capture it here
             accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
             capture_kwargs = {}
-            if accept_layer is not None and self.is_staged_run:
+            if accept_layer is not None:
                 capture_kwargs = {
                     "capture_layer_indices": [0, int(accept_layer)],
                     "return_hidden_states": True,
                 }
 
-            # Run thinker. Returns (text_hidden_states, captured_layer_dict)
-            # when capturing, a bare tensor otherwise.
-            return self.thinker(
+            # Run thinker
+            text_hidden_states, captured_layer_dict = self.thinker(
                 input_ids=input_ids,
                 positions=positions,
                 intermediate_tensors=intermediate_tensors,
@@ -425,6 +404,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 **capture_kwargs,
                 **kwargs,
             )
+            return text_hidden_states, captured_layer_dict
 
         # ========== Stage 2.1: Talker ==========
         elif self.model_stage == "talker":
@@ -514,11 +494,7 @@ class Qwen3OmniMoeForConditionalGeneration(
             return model_outputs
 
         if self.model_stage == "thinker":
-            if isinstance(model_outputs, tuple):
-                text_hidden_states, captured_layer_dict = model_outputs
-            else:
-                # Bare tensor: capture was not requested (no accept_hidden_layer).
-                text_hidden_states, captured_layer_dict = model_outputs, None
+            text_hidden_states, captured_layer_dict = model_outputs
             # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs.
             # These will later be projected into talker text space by the talker stage.
             multimodal_outputs: OmniPayload = captured_layer_dict if captured_layer_dict is not None else {}

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -1495,7 +1495,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         else:
             deepstack_input_embeds = None
 
-        model_output = self.language_model.model(
+        hidden_states, captured_hidden_states = self.language_model.model(
             input_ids,
             positions,
             intermediate_tensors,
@@ -1509,16 +1509,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         if inputs_embeds is not None and get_pp_group().is_first_rank:
             self._clear_deepstack_input_embeds(inputs_embeds.size(0))
 
-        if isinstance(model_output, IntermediateTensors):
-            # Non-last PP rank: forward the intermediate tensors as-is.
-            return model_output
-        hidden_states, captured_hidden_states = model_output
-        if capture_layer_indices is None and not return_hidden_states:
-            # No capture requested. Return a bare tensor: stock vLLM's
-            # GPUModelRunner indexes the forward output directly, so a tuple
-            # here breaks non-staged runs that resolve this class through the
-            # global registry (e.g. plain `vllm serve` of the thinker).
-            return hidden_states
         return hidden_states, captured_hidden_states
 
     def compute_logits(

vllm_omni/model_executor/models/utils.py

@@ -137,17 +137,3 @@ def reinit_rotary_inv_freq(
             inv_freq.copy_(new_inv_freq.to(dtype=inv_freq.dtype))
         n_fixed += 1
     return n_fixed
-
-
-def is_interleaved(config) -> bool:
-    """Detect if the model with this config is used with interleaved attention.
-
-    Replicates the helper that was removed from upstream
-    ``vllm.transformers_utils.config`` (vLLM commit 26d725c334,
-    "[Model] Add VaultGemma via Transformers modeling backend") so models
-    that still need the check can share one copy.
-    """
-    text_config = config.get_text_config()
-    if layer_types := getattr(text_config, "layer_types", None):
-        return len(set(layer_types)) > 1
-    return False

vllm_omni/model_extras/cosmos3.py

@@ -3,32 +3,6 @@
 
 from __future__ import annotations
 
-from typing import Any
-
-
-def build_text_to_image_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    height: int | None = None,
-    width: int | None = None,
-) -> dict[str, Any]:
-    """Build a Cosmos3 T2I prompt by selecting image output modality.
-
-    Cosmos3 uses the same pipeline class for T2I and video modes. The OpenAI
-    image endpoint selects T2I by adding ``modalities=["image"]``. ``height``
-    and ``width`` are accepted for registry compatibility; sizing travels
-    through sampling params instead of the prompt payload.
-    """
-    del height, width
-    text_prompt: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["image"],
-    }
-    if negative_prompt is not None:
-        text_prompt["negative_prompt"] = negative_prompt
-    return text_prompt
-
-
 COSMOS3_EXTRA_BODY_PARAMS = frozenset(
     {
         "flow_shift",

vllm_omni/model_extras/lingbot_video.py

@@ -3,11 +3,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
-from typing import Any
-
-from PIL import Image
-
 LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset(
     {
         "batch_cfg",
@@ -24,43 +19,3 @@ LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset(
         "t_thresh",
     }
 )
-
-
-def build_text_to_image_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    height: int | None = None,
-    width: int | None = None,
-) -> dict[str, Any]:
-    del height, width
-    result: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["image"],
-    }
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result
-
-
-def build_image_to_video_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    media_inputs: Mapping[str, Any],
-    height: int | None = None,
-    width: int | None = None,
-    num_frames: int | None = None,
-) -> dict[str, Any]:
-    del height, width
-    if set(media_inputs) != {"image"} or not isinstance(media_inputs["image"], Image.Image):
-        raise ValueError("LingBot text-image-to-video requires exactly one PIL image input.")
-
-    result: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["video"],
-        "multi_modal_data": {"image": media_inputs["image"]},
-    }
-    if num_frames is not None:
-        result["num_frames"] = num_frames
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result

vllm_omni/model_extras/registry.py

@@ -28,21 +28,12 @@ from vllm_omni.model_extras.cosmos3 import (
     COSMOS3_EXTRA_BODY_PARAMS,
     COSMOS3_EXTRA_OUTPUT_PARAMS,
 )
-from vllm_omni.model_extras.cosmos3 import (
-    build_text_to_image_prompt as build_cosmos3_text_to_image_prompt,
-)
 from vllm_omni.model_extras.helios import (
     HELIOS_EXTRA_BODY_PARAMS,
     HELIOS_EXTRA_OUTPUT_PARAMS,
 )
 from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt
 from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS
-from vllm_omni.model_extras.lingbot_video import (
-    build_image_to_video_prompt as build_lingbot_image_to_video_prompt,
-)
-from vllm_omni.model_extras.lingbot_video import (
-    build_text_to_image_prompt as build_lingbot_text_to_image_prompt,
-)
 from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS
 from vllm_omni.model_extras.magi_human import (
     MAGI_HUMAN_EXTRA_BODY_PARAMS,
@@ -147,18 +138,6 @@ def build_x_to_text_prompt(
     return builder(model, prompt, has_image)
 
 
-def default_text_to_image_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    height: int | None = None,
-    width: int | None = None,
-) -> dict[str, Any]:
-    result: dict[str, Any] = {"prompt": prompt}
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result
-
-
 def default_image_to_image_prompt(
     prompt: str,
     negative_prompt: str | None,
@@ -175,20 +154,6 @@ def default_image_to_image_prompt(
     return result
 
 
-def default_image_to_video_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    media_inputs: Mapping[str, Any],
-    height: int | None = None,
-    width: int | None = None,
-    num_frames: int | None = None,
-) -> dict[str, Any]:
-    del height, width, num_frames
-    if set(media_inputs) != {"image"} or not isinstance(media_inputs["image"], Image.Image):
-        raise ValueError("This model only supports a single --image input in the shared image-to-video example.")
-    return default_image_to_image_prompt(prompt, negative_prompt, media_inputs["image"])
-
-
 _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     "AudioXPipeline": {
         "extra_body_params": AUDIOX_EXTRA_BODY_PARAMS,
@@ -208,12 +173,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     "Cosmos3OmniDiffusersPipeline": {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
-        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
+        # The shared T2I example already supplies modalities=["image"].
     },
     "Cosmos3OmniPipeline": {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
-        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
     },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
@@ -229,8 +193,8 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     },
     "LingBotVideoPipeline": {
         "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
-        "text_to_image_prompt_builder": build_lingbot_text_to_image_prompt,
-        "image_to_video_prompt_builder": build_lingbot_image_to_video_prompt,
+        # Shared T2I/I2V envelopes select the output modality. LingBot's
+        # pipeline owns model-specific validation and normalization.
     },
     **{
         model_class_name: {
@@ -263,7 +227,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
 }
 
 # Multi-stage discovery reports the top-level wrapper rather than its DiT
-# submodule, so both names must resolve to the same request adapters.
+# submodule, so both names must resolve to the same request builders.
 _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
 _EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
 
@@ -305,18 +269,21 @@ def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None
 
 def build_text_to_image_prompt(
     model_class_name: str | None,
-    prompt: str,
-    negative_prompt: str | None,
+    prompt: dict[str, Any],
     height: int | None = None,
     width: int | None = None,
 ) -> dict[str, Any]:
+    """Build a model-specific T2I prompt from an example-owned envelope."""
     spec = _get_spec(model_class_name)
-    builder: TextToImagePromptBuilder = (
-        spec.get("text_to_image_prompt_builder", default_text_to_image_prompt)
-        if spec is not None
-        else default_text_to_image_prompt
+    builder: TextToImagePromptBuilder | None = spec.get("text_to_image_prompt_builder") if spec else None
+    if builder is None:
+        return prompt
+    return builder(
+        prompt=str(prompt["prompt"]),
+        negative_prompt=prompt.get("negative_prompt"),
+        height=height,
+        width=width,
     )
-    return builder(prompt, negative_prompt, height, width)
 
 
 def build_image_to_image_prompt(
@@ -338,17 +305,24 @@ def build_image_to_image_prompt(
 
 def build_image_to_video_prompt(
     model_class_name: str | None,
-    prompt: str,
-    negative_prompt: str | None,
-    media_inputs: Mapping[str, Any],
+    prompt: dict[str, Any],
     height: int | None = None,
     width: int | None = None,
     num_frames: int | None = None,
 ) -> dict[str, Any]:
+    """Build a model-specific I2V prompt from an example-owned envelope."""
     spec = _get_spec(model_class_name)
-    builder: ImageToVideoPromptBuilder = (
-        spec.get("image_to_video_prompt_builder", default_image_to_video_prompt)
-        if spec is not None
-        else default_image_to_video_prompt
+    builder: ImageToVideoPromptBuilder | None = spec.get("image_to_video_prompt_builder") if spec else None
+    if builder is None:
+        return prompt
+    media_inputs = prompt.get("multi_modal_data") or {}
+    if not isinstance(media_inputs, Mapping):
+        raise TypeError("Canonical I2V prompt multi_modal_data must be a mapping.")
+    return builder(
+        prompt=str(prompt["prompt"]),
+        negative_prompt=prompt.get("negative_prompt"),
+        media_inputs=media_inputs,
+        height=height,
+        width=width,
+        num_frames=num_frames,
     )
-    return builder(prompt, negative_prompt, media_inputs, height, width, num_frames)

vllm_omni/patch.py

@@ -419,93 +419,6 @@ def _patch_fp8_use_quack_fused_bias():
 _patch_fp8_use_quack_fused_bias()
 
 
-# =============================================================================
-# Patch torch inductor: prove factorable symbolic divisibility (CantSplit)
-# =============================================================================
-# WHY: torch 2.13's SizeVarAllocator.statically_known_multiple_of proves
-# symbolic divisibility via torch's own Mod (torch.utils._sympy.functions),
-# which deliberately stays unevaluated to bound compile time on wide
-# expressions. torch <= 2.11 used python `%` (sympy.Mod), whose eval factors
-# common terms. As a result an expression like
-#   15360*s31 + 15360*s87  vs  s31 + s87
-# — FLUX's fp8 graph with two dynamic sequence dims — is no longer provably
-# divisible, and inductor raises CantSplit on the first compiled forward,
-# killing the engine core during the warmup dummy run (nightly builds
-# 2953/2954, Diffusion Quantization Test).
-#
-# SCOPE: wraps statically_known_multiple_of. The wrapper can only ADD True
-# results, and only when sympy.cancel(numerator/denominator) yields a
-# polynomial with integer coefficients — then numerator == denominator * q
-# identically, so divisibility holds for every symbol assignment; this is
-# unconditionally sound. Everything else defers to the original result.
-# Guards: polynomial inputs only (inductor's FloorDiv/ModularIndexing atoms
-# never reach sympy.cancel) and <= 20 free symbols (mirrors upstream's own
-# cost cap).
-#
-# FRAGILITY / REMOVE WHEN: self-extinguishes at install time by probing
-# whether the unpatched method already proves the canonical factorable case
-# (as torch <= 2.11 does). When upstream restores factor-aware proving, the
-# probe passes and nothing is patched.
-def _provably_factorable_multiple(numerator, denominator) -> bool:
-    """True only when numerator/denominator cancels to an integer polynomial."""
-    try:
-        import sympy
-
-        num = sympy.sympify(numerator)
-        den = sympy.sympify(denominator)
-        if den.is_zero:
-            return False
-        symbols = num.free_symbols | den.free_symbols
-        if len(symbols) > 20:
-            return False
-        if not (num.is_polynomial(*symbols) and den.is_polynomial(*symbols)):
-            return False
-        quotient = sympy.cancel(num / den)
-        _, q_den = sympy.fraction(sympy.together(quotient))
-        if q_den != 1:
-            return False
-        if quotient.is_Integer:
-            return True
-        poly = sympy.Poly(quotient, *sorted(quotient.free_symbols, key=str))
-        return all(coeff.is_integer for coeff in poly.coeffs())
-    except Exception:  # noqa: BLE001 - a proof failure must never break compile
-        return False
-
-
-def _patch_inductor_factorable_divisibility():
-    try:
-        import sympy
-        from torch._inductor.sizevars import SizeVarAllocator
-    except ImportError:
-        return
-
-    original = SizeVarAllocator.statically_known_multiple_of
-    if getattr(original, "_vllm_omni_factorable_divisibility", False):
-        return
-
-    try:
-        _a, _b = sympy.symbols("_omni_probe_a _omni_probe_b", positive=True, integer=True)
-        if original(SizeVarAllocator(), 7 * _a + 7 * _b, _a + _b):
-            _PATCH_LOGGER.info("inductor factorable-divisibility patch: skipped (upstream already proves it).")
-            return
-    except Exception:  # noqa: BLE001
-        # Probe broke (constructor drift etc.) — install anyway; the wrapper
-        # never subtracts results from the original.
-        pass
-
-    def statically_known_multiple_of(self, numerator, denominator):
-        if original(self, numerator, denominator):
-            return True
-        return _provably_factorable_multiple(numerator, denominator)
-
-    statically_known_multiple_of._vllm_omni_factorable_divisibility = True
-    SizeVarAllocator.statically_known_multiple_of = statically_known_multiple_of
-    _PATCH_LOGGER.info("inductor factorable-divisibility patch: installed.")
-
-
-_patch_inductor_factorable_divisibility()
-
-
 # =============================================================================
 # Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown
 # =============================================================================

vllm_omni/platforms/musa/platform.py

@@ -3,8 +3,6 @@
 
 
 import torch
-from vllm.config import VllmConfig
-from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.interface import DeviceCapability
 from vllm_musa.platform import MUSAPlatformBase
@@ -32,6 +30,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     def get_omni_generation_worker_cls(cls) -> str:
         return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
 
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/deploy"
+
     @classmethod
     def has_flash_attn_package(cls) -> bool:
         from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed
@@ -113,10 +115,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         """MUSA supports torch.compile with inductor backend."""
         return True
 
-    @classmethod
-    def get_default_stage_config_path(cls) -> str:
-        return "vllm_omni/deploy"
-
     @classmethod
     def supports_talker_mtp_graph_capture(cls) -> bool:
         """MUSA keeps Qwen3 talker MTP outside its dedicated FULL graph."""
@@ -144,7 +142,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
         """Get the compute capability of the MUSA device."""
-        major, minor = torch.musa.get_device_capability(device_id)
+        major, minor = torch.cuda.get_device_capability(device_id)
         return DeviceCapability(major=major, minor=minor)
 
     @classmethod
@@ -185,26 +183,3 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.musa.get_device_name(device_id)
-
-    @classmethod
-    def set_device_control_env_var(cls, devices: str | int | None) -> None:
-        import os
-
-        os.environ["MUSA_VISIBLE_DEVICES"] = devices
-
-    @classmethod
-    def unset_device_control_env_var(cls) -> None:
-        import os
-
-        os.environ.pop("MUSA_VISIBLE_DEVICES", None)
-
-    @classmethod
-    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Prefer native ops while compiling, otherwise prefer vLLM kernels."""
-        from vllm.config.compilation import CompilationMode
-
-        cc = vllm_config.compilation_config
-        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["vllm_c", "native"]
-
-        return IrOpPriorityConfig.with_default(default)

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -11,7 +11,6 @@ from copy import copy, deepcopy
 
 import numpy as np
 import torch
-import vllm.v1.worker.mamba_utils as mamba_utils
 from vllm.config import CUDAGraphMode
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
@@ -23,6 +22,7 @@ from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs

vllm_omni/platforms/xpu/platform.py

@@ -130,12 +130,12 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
         """Copied from upstream XPUPlatform with inductor-aware logic.
 
         When inductor is active (compiling) use native as the default;
-        otherwise prefer vllm_c where available.
+        otherwise prefer xpu_kernels where available.
         """
         from vllm.config.compilation import CompilationMode
 
         cc = vllm_config.compilation_config
         using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["vllm_c", "native"]
+        default = ["native"] if using_inductor else ["xpu_kernels", "native"]
 
         return IrOpPriorityConfig.with_default(default)

vllm_omni/quantization/_copy_missing_attrs.py

@@ -1,24 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Shared weight-attribute helper for vLLM-omni quantization configs.
-
-``copy_missing_attrs`` used to live in upstream
-``vllm.model_executor.layers.quantization.fp8`` but was removed from vLLM
-(commit b5bcb3ce88, "[Refactor] Remove dead code in multiple files").
-vLLM-omni's diffusion quant methods still rely on it, so it is re-homed here.
-"""
-
-from typing import Any
-
-import torch
-from vllm.model_executor.utils import set_weight_attrs
-
-
-def copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None:
-    """Copies any attrs present in `old` but not in `new` to `new`."""
-    new_attrs = set(dir(new))
-    attrs_to_set: dict[str, Any] = {}
-    for attr in dir(old):
-        if attr not in new_attrs:
-            attrs_to_set[attr] = getattr(old, attr)
-    set_weight_attrs(new, attrs_to_set)

vllm_omni/quantization/bitsandbytes_config.py

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     is_layer_skipped,
 )
@@ -28,9 +29,6 @@ from vllm.model_executor.parameter import ModelWeightParameter
 from vllm.model_executor.utils import replace_parameter
 
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.quantization._copy_missing_attrs import (
-    copy_missing_attrs as _copy_missing_attrs,
-)
 from vllm_omni.quantization.int8_config import LazyWeightMixin
 
 if TYPE_CHECKING:

vllm_omni/quantization/int8_config.py

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter
+from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter, _copy_missing_attrs
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     is_layer_skipped,
 )
@@ -38,9 +38,6 @@ from vllm.model_executor.parameter import (
 from vllm.model_executor.utils import replace_parameter
 
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.quantization._copy_missing_attrs import (
-    copy_missing_attrs as _copy_missing_attrs,
-)
 
 if current_omni_platform.is_npu():
     import torch_npu

vllm_omni/quantization/mxfp4_config.py

@@ -66,15 +66,13 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs
 from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
 from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
 from vllm.model_executor.parameter import ModelWeightParameter
 from vllm.model_executor.utils import replace_parameter
 
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.quantization._copy_missing_attrs import (
-    copy_missing_attrs as _copy_missing_attrs,
-)
 from vllm_omni.quantization.mxfp8_config import (
     MXFPLinearMethodBase,
     _LazyWeightMixin,

vllm_omni/quantization/mxfp8_config.py

@@ -49,16 +49,13 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter
+from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter, _copy_missing_attrs
 from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
 from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
 from vllm.model_executor.parameter import ModelWeightParameter
 from vllm.model_executor.utils import replace_parameter
 
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.quantization._copy_missing_attrs import (
-    copy_missing_attrs as _copy_missing_attrs,
-)
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.utils import WeightsMapper

vllm_omni/worker/base.py

@@ -121,11 +121,6 @@ class OmniGPUWorkerBase(GPUWorker):
 
         self.non_torch_memory = profile_result.non_torch_increase
         self.peak_activation_memory = profile_result.torch_peak_increase
-        # Upstream 58b2012aa2 added `total_consumed` to the profiling result
-        # and reads it in GPUWorker.compile_or_warm_up_model() (when
-        # kv_cache_memory_bytes is None and peak_activation_memory is set, both
-        # true here). Mirror upstream so the omni override keeps it populated.
-        self.total_consumed = profile_result.total_consumed
 
         process_memory = (
             get_process_gpu_memory(self.local_rank)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -577,7 +577,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
         profile_seq_lens: int | None = None,
-        randomize_inputs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -837,13 +836,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     num_tokens_across_dp[:] = num_tokens_padded
 
             with (
-                # Never randomize generation-stage dummy ids: they are
-                # structured payloads (e.g. packed RVQ codec codes that get
-                # code_offset added before embedding), so vocab-uniform random
-                # ids index past the codebook tables — CUDA device-side assert
-                # in code2wav warmup. The expert-balance rationale behind
-                # randomize_inputs targets text-MoE models, which generation
-                # stages are not. The kwarg stays accepted for vLLM API compat.
                 self.maybe_randomize_inputs(input_ids, inputs_embeds),
                 set_forward_context(
                     attn_metadata,

vllm_omni/worker/gpu_model_runner.py

@@ -882,7 +882,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         is_graph_capturing: bool = False,
         num_active_loras: int = 0,
         profile_seq_lens: int | None = None,
-        randomize_inputs: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -910,9 +909,6 @@ class OmniGPUModelRunner(GPUModelRunner):
             profile_seq_lens: If provided, use this value for seq_lens instead
                 of max_query_len. Used to profile attention workspace that
                 scales with context length.
-            randomize_inputs: If True, randomize dummy input ids to balance
-                expert selection. vLLM's kernel_warmup passes this when
-                autotuning flashinfer.
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -1135,7 +1131,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     num_tokens_across_dp[:] = num_tokens_padded
 
             with (
-                self.maybe_randomize_inputs(input_ids, inputs_embeds, randomize_inputs=randomize_inputs),
+                self.maybe_randomize_inputs(input_ids, inputs_embeds),
                 set_forward_context(
                     attn_metadata,
                     self.vllm_config,