[Dependency] Upgrade Cache-DiT to 1.5.0

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-nightly.yml

@@ -345,7 +345,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json

.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:

.buildkite/npu/test-npu-nightly.yml

@@ -115,7 +115,7 @@ steps:
           DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

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

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -82,7 +82,7 @@ example and the DiT sections below.
 - Optional environment variables:
 
 ```bash
-export CACHE_DIT_VERSION=1.3.0
+export CACHE_DIT_VERSION=1.5.0
 ```
 
 HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`

requirements/common.txt

@@ -16,7 +16,7 @@ diffusers==0.38.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
-cache-dit==1.3.0
+cache-dit==1.5.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625

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/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/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/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/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/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/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/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/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

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},
                     )

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/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/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/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,