[Rebase] Rebase to vllm 0.28.0 (#6606)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

docker/Dockerfile.ci

@@ -1,4 +1,4 @@
-# Target vllm version: v0.27.0 (releases/v0.27.0 branch).
+# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
 # 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
@@ -14,11 +14,14 @@
 # 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.
+# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
+# rode on the v0.27.0 base image with the full reinstall block until the
+# official v0.28.0 image shipped, at which point this file dropped back to the
+# form you are reading. 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.28.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents

docker/Dockerfile.cuda

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

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
 
 vLLM-Omni supports the following hardware platforms:
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
 
 ```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
 ```
 
-The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.26.0
+git checkout v0.28.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.26.0 \
+    vllm/vllm-omni:v0.28.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
   -t vllm-omni-cuda .
 ```
 

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
 
 ```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
 ```
 
-The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.26.0
+git checkout v0.28.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.26.0 \
+  vllm/vllm-omni-rocm:v0.28.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.26.0
+  vllm/vllm-omni-rocm:v0.28.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

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

examples/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""vLLM-Omni example scripts, importable as a package.
+
+Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
+places an ``examples`` package on ``sys.path``), a regular package at this repo
+root wins over the upstream ``examples`` namespace during ``import examples``.
+"""

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
     return sched

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -93,6 +93,7 @@ def _run_resumable_segment_stop(
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 
@@ -227,6 +228,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.pending_stop_after_extraction = set()
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     sched.recompute_kv_load_failures = False
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
 

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
+
+Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
+non-empty scheduled step and, at the top of ``update_from_output()``,
+advances ``processed_step_seq`` and drains the deferred frees. The Omni
+schedulers each carried only half of that contract:
+
+- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
+  advances) but its reimplemented ``update_from_output()`` never advanced
+  ``processed_step_seq`` nor drained -> deferred blocks accumulate without
+  bound under async/PP KV-consumer configurations.
+- ``OmniGenerationScheduler`` reimplements both: the fast-path
+  ``schedule()`` never advanced ``sched_step_seq`` (so finished requests
+  freed blocks an in-flight GPU step could still be writing) and its
+  ``update_from_output()`` never drained.
+
+These tests pin the restored contract at the call boundary; the drain and
+free logic itself is upstream-tested.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus 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.interface import PauseState
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import SchedulingPolicy
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_update_sched(scheduler_cls) -> MagicMock:
+    """MagicMock scheduler driving the real ``update_from_output`` with an
+    empty batch; only the fence state carries real values."""
+    sched = MagicMock()
+    sched.requests = {}
+    sched.perf_metrics = None
+    sched.chunk_transfer_adapter = None
+    sched.defer_block_free = True
+    sched.processed_step_seq = 0
+    return sched
+
+
+def _empty_frame(*, total_num_scheduled_tokens: int):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {}
+    scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = []
+    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
+    return scheduler_output, model_runner_output
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
+    sched = _make_update_sched(scheduler_cls)
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 1
+    sched._drain_deferred_frees.assert_called_once_with()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
+    """0-token steps do not advance the fence upstream (their seq was never
+    advanced in schedule()); draining on them would run ahead of the GPU."""
+    sched = _make_update_sched(scheduler_cls)
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 0
+    sched._drain_deferred_frees.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
+    sched = _make_update_sched(scheduler_cls)
+    sched.defer_block_free = False
+    scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+    scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert sched.processed_step_seq == 0
+    sched._drain_deferred_frees.assert_not_called()
+
+
+def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
+    """MagicMock scheduler that drives the real generation fast-path
+    ``schedule()`` over one running one-shot request."""
+    request = Request(
+        request_id="req-fence",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+
+    sched = MagicMock()
+    sched.max_num_scheduled_tokens = 8
+    sched._pause_state = PauseState.UNPAUSED
+    sched.requests = {request.request_id: request}
+    sched.running = [request]
+    sched.waiting = []  # falsy: the waiting loop never runs
+    sched.policy = SchedulingPolicy.FCFS
+    sched.chunk_transfer_adapter = None
+    sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
+    sched.num_lookahead_tokens = 0
+    sched.log_stats = False
+    sched.kv_cache_config = MagicMock(kv_cache_groups=[])
+    sched.use_v2_model_runner = False
+    sched.prev_step_scheduled_req_ids = set()
+    sched.needs_kv_cache_zeroing = False
+    sched.connector = None
+    sched.ec_connector = None
+    sched.defer_block_free = defer_block_free
+    sched.sched_step_seq = 0
+    # last_sched_seq stamping order: _update_after_schedule must observe the
+    # already-advanced fence value, as upstream schedule() guarantees.
+    sched.seq_at_update_after_schedule = None
+
+    def _capture(_scheduler_output):
+        sched.seq_at_update_after_schedule = sched.sched_step_seq
+
+    sched._update_after_schedule.side_effect = _capture
+    return sched, request
+
+
+def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
+    sched, request = _make_fast_path_sched(defer_block_free=True)
+
+    OmniGenerationScheduler.schedule(sched)
+
+    assert sched.sched_step_seq == 1
+    assert sched.seq_at_update_after_schedule == 1
+
+
+def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
+    sched, request = _make_fast_path_sched(defer_block_free=False)
+
+    OmniGenerationScheduler.schedule(sched)
+
+    assert sched.sched_step_seq == 0

tests/core/sched/test_omni_sched_ec_request_finish.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
+
+Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
+BEFORE the encoder cache is freed (so the connector can inspect per-request
+state), honors its delayed-free result, and returns the generated
+``ec_transfer_params``; ``update_from_output()`` then carries them into the
+``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
+always returned ``None``; both Omni schedulers discarded the second tuple
+item, so a remote encoder-cache handle could never reach the frontend.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus 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
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
+
+
+def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
+    """Minimal real-`_free_request` harness; `order` records the EC hook
+    firing relative to the encoder-cache free."""
+    order = MagicMock()
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._omits_kv_transfer_cache = {}
+    sched._connector_finished = lambda request: (False, None)
+    sched.ec_connector = ec_connector
+    if ec_connector is not None:
+        order.attach_mock(ec_connector.request_finished, "request_finished")
+    sched.encoder_cache_manager = MagicMock()
+    order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
+    sched.finished_req_ids = set()
+    sched._new_prompt_len_snapshot = {}
+    sched.finished_req_ids_dict = None
+    sched._should_transfer_kv_for_request = lambda req_id: False
+    sched._free_blocks = MagicMock()
+    sched._free_input_coordinator_request = MagicMock()
+    sched.chunk_transfer_adapter = None
+    return sched, order
+
+
+class _FakeFinishedRequest:
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+
+    def is_finished(self) -> bool:
+        return True
+
+
+def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
+    ec_connector = MagicMock()
+    ec_connector.request_finished.return_value = (False, EC_PARAMS)
+    sched, order = _make_free_request_sched(ec_connector=ec_connector)
+    request = _FakeFinishedRequest("req-ec")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params == EC_PARAMS
+    ec_connector.request_finished.assert_called_once_with(request)
+    hook_calls = [name for name, *_ in order.mock_calls]
+    assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
+    sched._free_blocks.assert_called_once_with(request)
+
+
+def test_free_request_honors_ec_delayed_free() -> None:
+    ec_connector = MagicMock()
+    ec_connector.request_finished.return_value = (True, EC_PARAMS)
+    sched, _ = _make_free_request_sched(ec_connector=ec_connector)
+    request = _FakeFinishedRequest("req-ec-delay")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params == EC_PARAMS
+    sched._free_blocks.assert_not_called()
+
+
+def test_free_request_without_ec_connector_returns_none() -> None:
+    sched, _ = _make_free_request_sched(ec_connector=None)
+    request = _FakeFinishedRequest("req-no-ec")
+
+    kv_params, ec_params = sched._free_request(request)
+
+    assert ec_params is None
+    sched._free_blocks.assert_called_once_with(request)
+
+
+def _make_finishing_request() -> Request:
+    request = Request(
+        request_id="req-ec-finish",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+    return request
+
+
+def _make_finish_sched(request: Request) -> MagicMock:
+    """MagicMock scheduler driving the real ``update_from_output`` through a
+    finishing request; ``_free_request`` reports EC transfer params."""
+    sched = MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._maybe_decode_pooling_output.return_value = None
+    sched._handle_stopped_request.return_value = True
+    sched._free_request.return_value = (None, EC_PARAMS)
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.chunk_transfer_adapter = None
+    sched.finished_req_ids_dict = None
+    sched._new_prompt_len_snapshot = {}
+    return sched
+
+
+def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
+    request = _make_finishing_request()
+    sched = _make_finish_sched(request)
+    sched._update_request_with_output.return_value = ([42], True)
+
+    engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.ec_transfer_params == EC_PARAMS
+    sched._free_request.assert_called_once_with(request)
+
+
+def _finish_frame_parts(request: Request):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    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 = [[42]]
+    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 = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    return scheduler_output, model_runner_output
+
+
+def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
+    request = _make_finishing_request()
+    # One-shot generation finish: the whole prompt is computed.
+    request.num_computed_tokens = len(request.prompt_token_ids)
+    sched = _make_finish_sched(request)
+
+    engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.ec_transfer_params == EC_PARAMS
+    sched._free_request.assert_called_once_with(request)

tests/core/sched/test_omni_sched_prefill_stats_finalize.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
+
+Upstream v0.28 finalizes prefill statistics with
+``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
+before emitting them, which computes ``num_cache_creation_tokens`` (tokens
+computed and written to the prefix cache). Both Omni schedulers previously
+emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus 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
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_request() -> Request:
+    request = Request(
+        request_id="req-prefill-stats",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=4),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+    request.status = RequestStatus.RUNNING
+    # A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
+    request.prefill_stats.set(
+        num_prompt_tokens=3,
+        num_local_cached_tokens=1,
+        num_external_cached_tokens=0,
+    )
+    return request
+
+
+def _make_sched(request: Request) -> MagicMock:
+    sched = MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._maybe_decode_pooling_output.return_value = None
+    sched._handle_stopped_request.return_value = True
+    # The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
+
+    def _free_and_invalidate(req):
+        # Lifecycle guard: freeing releases the KV blocks, after which the
+        # manager reports 0 cached tokens. If finalize ran after
+        # _free_request (the reviewed bug), creation count would compute
+        # against 0 and the assertions below would see 0, not 2.
+        sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+        return (None, None)
+
+    sched._free_request.side_effect = _free_and_invalidate
+    sched.chunk_transfer_adapter = None
+    sched.finished_req_ids_dict = None
+    sched._new_prompt_len_snapshot = {}
+    return sched
+
+
+def _frame(request: Request):
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    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 = [[42]]
+    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 = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    return scheduler_output, model_runner_output
+
+
+def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+    request = _make_request()
+    sched = _make_sched(request)
+    sched._update_request_with_output.return_value = ([42], True)
+
+    engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.prefill_stats is not None
+    assert eco.prefill_stats.num_cached_tokens == 1
+    assert eco.prefill_stats.num_cache_creation_tokens == 2
+    sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
+
+
+def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+    request = _make_request()
+    request.num_computed_tokens = len(request.prompt_token_ids)
+    sched = _make_sched(request)
+
+    engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
+
+    (eco,) = engine_core_outputs[request.client_index].outputs
+    assert eco.prefill_stats is not None
+    assert eco.prefill_stats.num_cached_tokens == 1
+    assert eco.prefill_stats.num_cache_creation_tokens == 2
+    sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Layer-level correctness & performance test for MoT parallel linear layers.
 
@@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
     return output
 
 
+def _assert_gen_expert_parallel_metadata(layer):
+    """The secondary expert must satisfy vLLM's parallel-linear contract."""
+    for attr in (
+        "tp_rank",
+        "tp_size",
+        "input_size",
+        "input_size_per_partition",
+        "output_size",
+        "output_size_per_partition",
+    ):
+        assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
+
+
+def test_mot_gen_expert_parallel_metadata():
+    """Keep secondary experts compatible with vLLM online quantizers."""
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        _is_tp_sharded,
+    )
+
+    with set_current_vllm_config(VllmConfig()):
+        layers = (
+            MoTQKVParallelLinear(
+                hidden_size=16,
+                head_size=4,
+                total_num_heads=4,
+                total_num_kv_heads=2,
+                bias=False,
+                vae_bias=False,
+                params_dtype=torch.bfloat16,
+                disable_tp=True,
+            ),
+            MoTRowParallelLinear(
+                16,
+                8,
+                bias=False,
+                vae_bias=False,
+                input_is_parallel=True,
+                params_dtype=torch.bfloat16,
+                disable_tp=True,
+            ),
+        )
+
+    for layer in layers:
+        _assert_gen_expert_parallel_metadata(layer)
+        assert _is_tp_sharded(layer.gen_exp) is False
+
+
 def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
     """Compare outputs, print metrics, and assert correctness.
 
@@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
             disable_tp=True,
         ).cuda()
 
+        _assert_gen_expert_parallel_metadata(mot_linear)
+
         assert text_linear.output_size_per_partition == N, (
             f"Expected output_size_per_partition={N}, "
             f"got {text_linear.output_size_per_partition}. "
@@ -423,6 +475,8 @@ def test_mot_o_proj(
             disable_tp=True,
         ).cuda()
 
+        _assert_gen_expert_parallel_metadata(mot_linear)
+
         _sync_weights(text_linear, vae_linear, mot_linear)
 
         text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
@@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
         ).cuda()
 
         with torch.no_grad():
+            # vLLM linear layers leave weights uninitialized (torch.empty), so
+            # a freshly zeroed GPU page would make both outputs zero and the
+            # cosine-similarity gate degenerate. Seed deterministic values so
+            # und-mode is compared against a meaningful reference.
+            ref_linear.weight.normal_(mean=0.0, std=0.02)
+            if ref_linear.bias is not None:
+                ref_linear.bias.normal_(mean=0.0, std=0.02)
             mot_linear.weight.copy_(ref_linear.weight)
             if bias and ref_linear.bias is not None:
                 mot_linear.bias.copy_(ref_linear.bias)
@@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
         ).cuda()
 
         with torch.no_grad():
+            # vLLM linear layers leave weights uninitialized (torch.empty), so
+            # a freshly zeroed GPU page would make both outputs zero and the
+            # cosine-similarity gate degenerate. Seed deterministic values so
+            # und-mode is compared against a meaningful reference.
+            ref_linear.weight.normal_(mean=0.0, std=0.02)
+            if ref_linear.bias is not None:
+                ref_linear.bias.normal_(mean=0.0, std=0.02)
             mot_linear.weight.copy_(ref_linear.weight)
             if bias and ref_linear.bias is not None:
                 mot_linear.bias.copy_(ref_linear.bias)

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -245,7 +245,7 @@ else:
         [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
         capture_output=True,
         text=True,
-        timeout=3,
+        timeout=30,
         check=False,
     )
 

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -1568,6 +1568,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     scheduler.structured_output_manager.should_advance.return_value = False
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
 
     request = _HashableRequest(
@@ -1648,6 +1649,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager = mocker.MagicMock()
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
     scheduler.waiting_for_transfer_free = set()
     scheduler.transfer_triggered_requests = set()
@@ -1803,6 +1805,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     scheduler.structured_output_manager.should_advance.return_value = False
     scheduler.finished_req_ids_dict = {}
     scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
     scheduler.kv_event_publisher = mocker.MagicMock()
     scheduler._pending_finish_reqs = list(pending_finish_reqs)
 

tests/e2e/offline_inference/test_flux2_klein.py

@@ -76,6 +76,7 @@ def _send_inpaint_with_generator(
 
 @pytest.mark.slow
 @pytest.mark.diffusion
+@pytest.mark.advanced_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
     input_image, mask_image = _create_test_inputs()
@@ -97,6 +98,7 @@ def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
 
 @pytest.mark.slow
 @pytest.mark.diffusion
+@pytest.mark.advanced_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_flux2_klein_inpaint_different_seeds_different_output(offline_client: OfflineOmniClient):
     input_image, mask_image = _create_test_inputs()

tests/engine/test_async_omni_engine_stage_init.py

@@ -452,6 +452,10 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
         "env": {runtime_env_var: "stage-value"},
     }
     monkeypatch.setenv(runtime_env_var, "parent-value")
+    # Pin the logical->physical device mapping so the assertion below is
+    # deterministic regardless of the ambient CUDA_VISIBLE_DEVICES (which on a
+    # rebase host can be e.g. "3,4", mapping logical device 0 -> physical 3).
+    runtime._init_visible_devices_baseline = "0"
 
     captured: dict[str, str | None] = {}
 

tests/engine/test_output_processor.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
 
 from types import SimpleNamespace
@@ -581,18 +581,18 @@ def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
 
 def test_no_detokenizer_make_request_output_with_routed_experts():
     """make_request_output accepts the routed_experts arg that the multimodal
-    output channel (_process_mm_only_outputs) passes positionally, and attaches
-    it to the completion output.
+    output channel (_process_mm_only_outputs) passes, and attaches it to the
+    completion output.
 
-    Regression: the call site passes 6 positional args
-    (..., kv_transfer_params, routed_experts) but the override previously took
-    only 5, raising TypeError on every generation-stage output.
+    routed_experts is omni-specific keyword-only: the 6th positional now
+    matches upstream's ec_transfer_params so that super().process_outputs()
+    calls do not misroute it.
     """
     s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
     s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
     routed_experts = np.zeros((2, 3), dtype=np.int32)
     # Mirror the exact call shape of _process_mm_only_outputs.
-    result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts)
+    result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts)
     assert result is not None
     assert not isinstance(result, PoolingRequestOutput)
     assert result.outputs[0].routed_experts is routed_experts
@@ -663,3 +663,40 @@ def test_mm_only_outputs_update_iteration_stats():
     assert finished.finish_reason == FinishReason.STOP
     assert finished.num_prompt_tokens == state.prompt_len
     assert finished.num_generation_tokens == 2
+
+
+def test_ec_transfer_params_survive_both_construction_paths():
+    """v0.28 contract: encoder-cache transfer metadata must reach
+    RequestOutput.ec_transfer_params through BOTH omni construction paths —
+    the upstream-delegating one (logprobs processor present) and the
+    no-detokenizer generation-stage one (direct RequestOutput build).
+    Regression: the override accepted the argument but dropped it."""
+    ec = {"remote_handle": "enc-cache-1"}
+
+    state = _make_state(RequestOutputKind.CUMULATIVE)
+    out = state.make_request_output([1], None, None, None, {"kv": 1}, ec)
+    assert out is not None and out.ec_transfer_params == ec
+
+    kwargs = dict(_DEFAULT_STATE_KWARGS)
+    kwargs.update(logprobs_processor=None, detokenizer=None)
+    gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
+    completion = gen_state._new_completion_output([1], None, None)
+    direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec)
+    assert direct.ec_transfer_params == ec
+
+
+def test_num_cache_creation_tokens_reaches_direct_request_output():
+    """v0.28 contract: prefix-cache creation usage must reach
+    RequestOutput.num_cache_creation_tokens through the no-detokenizer
+    direct build, which previously passed only num_cached_tokens."""
+    kwargs = dict(_DEFAULT_STATE_KWARGS)
+    kwargs.update(logprobs_processor=None, detokenizer=None)
+    gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
+    gen_state.num_cached_tokens = 3
+    gen_state.num_cache_creation_tokens = 2
+
+    completion = gen_state._new_completion_output([1], None, None)
+    direct = gen_state._new_request_output("r", [completion], False, None, None)
+
+    assert direct.num_cached_tokens == 3
+    assert direct.num_cache_creation_tokens == 2

tests/engine/test_single_stage_mode.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer."""
 
 from __future__ import annotations
@@ -1374,7 +1377,7 @@ class TestConnectRemoteEngineCoresCoordinator:
         omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2)
         omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2)
         mock_wait.assert_called_once()
-        _, _, core_engines, parallel_config, *_ = mock_wait.call_args.args
+        _, core_engines, parallel_config, *_ = mock_wait.call_args.args
         assert core_engines[0].local is False
         assert parallel_config.data_parallel_size_local == 0
 

tests/entrypoints/test_omni_snapshot_download.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """``omni_snapshot_download`` must honor vLLM's ``VLLM_USE_MODELSCOPE`` semantics.
 
-vLLM treats the flag as enabled only for the literal string ``"true"``
-(case-insensitive). Reading ``os.environ`` directly made every non-empty value
-truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0`` still took the
-ModelScope path.
+vLLM treats the flag as enabled for the literal strings ``"1"`` or ``"true"``
+(case-insensitive; see vllm.envs). Reading ``os.environ`` directly made every
+non-empty value truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0``
+still took the ModelScope path.
 """
 
 import sys
@@ -127,7 +127,7 @@ def download_backend(monkeypatch: pytest.MonkeyPatch):
     return run
 
 
-@pytest.mark.parametrize("value", ["0", "1", "False", "false", "no", "off"])
+@pytest.mark.parametrize("value", ["0", "False", "false", "no", "off"])
 def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend, value):
     monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
 
@@ -135,7 +135,7 @@ def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend,
     assert download_backend() == "huggingface"
 
 
-@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
+@pytest.mark.parametrize("value", ["1", "true", "True", "TRUE"])
 def test_true_values_enable_modelscope(monkeypatch, download_backend, value):
     monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
 

tests/worker/test_gpu_ar_model_runner.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import copy
 import re
@@ -894,6 +897,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
         lambda *args, **kwargs: (
             0,
             None,
+            None,
             [],
             None,
             ["r1", "r2"],

vllm_omni/config/endpoint_policy.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Endpoint restriction policy for omni pipelines."""
 
 from dataclasses import dataclass
@@ -8,7 +8,7 @@ from typing import NamedTuple
 
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-from vllm.entrypoints.serve.utils.error_response import create_error_response
+from vllm.entrypoints.serve import create_error_response
 
 
 class RouteTarget(NamedTuple):

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -344,8 +344,25 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None)
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
+        ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
         cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
 
+        # Every GPU write enqueued by this and earlier steps has completed, so
+        # it is safe to return deferred-free blocks to the pool. This is the
+        # update-side half of the upstream v0.28 deferred-free fence; the schedule
+        # half advances sched_step_seq inside super().schedule().
+        # getattr: __new__-constructed test schedulers carry no
+        # defer_block_free attribute.
+        if (
+            getattr(self, "defer_block_free", False)
+            # getattr again: SchedulerOutput is mocked field-by-field in the
+            # scheduler unit tests, and dataclass fields without defaults are
+            # invisible to MagicMock(spec=...).
+            and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0
+        ):
+            self.processed_step_seq += 1
+            self._drain_deferred_frees()
+
         perf_stats: PerfStats | None = None
         if self.perf_metrics and self.perf_metrics.is_enabled():
             perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
@@ -491,6 +508,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             mm_output = mm_outputs[req_index] if mm_outputs else None
             inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None
             kv_transfer_params = None
+            ec_transfer_params = None
             finish_reason = None
             routed_experts = None
 
@@ -542,6 +560,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     request.resumable = False
                     stopped = True
 
+            # Finalize prefill stats BEFORE stop handling (upstream v0.28
+            # order): _free_request below releases the KV blocks, after which
+            # estimate_cached_tokens(request) reports 0. kv_transfer_params is
+            # omitted from the emission predicate here because it only becomes
+            # non-None when stopped is already True.
+            prefill_stats = None
+            if new_token_ids or mm_output is not None or pooler_output is not None or stopped:
+                prefill_stats = request.take_prefill_stats()
+                if prefill_stats is not None:
+                    prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request))
+
             confirmed_num_computed_tokens = None
             if stopped:
                 if self.chunk_transfer_adapter is not None:
@@ -597,7 +626,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     request.spec_token_ids = []
                     request._output_token_ids.clear()
                 if finished:
-                    kv_transfer_params, _ = self._free_request(request)
+                    kv_transfer_params, ec_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
                 elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
@@ -625,8 +654,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
-                    prefill_stats=request.take_prefill_stats(),
+                    prefill_stats=prefill_stats,
                     kv_transfer_params=kv_transfer_params,
+                    ec_transfer_params=ec_transfer_params,
                     routed_experts=routed_experts,
                     num_nans_in_logits=request.num_nans_in_logits,
                     is_segment_finished=is_segment_finished,
@@ -666,6 +696,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
+        # EC Connector: update state from worker-side EC connector output.
+        # Use getattr for safety with test __new__/SimpleNamespace code paths.
+        if getattr(self, "ec_connector", None) is not None and ec_connector_output:
+            self.ec_connector.update_connector_output(ec_connector_output)
+
         kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
         self._publish_kv_cache_events()
 
@@ -795,6 +830,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # 1. Standard cleanup parts from base _free_request
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
 
+        # EC Connector: mirror the KV hook (upstream v0.28 _free_request).
+        # The contract requires firing before the encoder cache is freed so
+        # the connector can inspect per-request state (e.g. which mm_hashes
+        # it recorded during save_caches()) and emit ec_transfer_params for
+        # the response body. getattr: __new__-constructed test schedulers
+        # carry no ec_connector attribute.
+        ec_xfer_params: dict[str, Any] | None = None
+        if getattr(self, "ec_connector", None) is not None:
+            ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request)
+            connector_delay_free_blocks |= ec_delay_free
+
         self.encoder_cache_manager.free(request)
         request_id = request.request_id
         self.finished_req_ids.add(request_id)
@@ -820,10 +866,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         self.waiting_for_transfer_free.add(request_id)
                         self._kv_wait_start_ts[request_id] = time.monotonic()
                         kv_xfer_params = None
-                        return kv_xfer_params, None
+                        return kv_xfer_params, ec_xfer_params
                     elif request_id in self.waiting_for_transfer_free:
                         # Blocks held until KV extraction completes in a future step.
-                        return None, None
+                        return None, ec_xfer_params
                     else:
                         logger.debug(
                             f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
@@ -861,14 +907,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         if isinstance(add_info, dict):
                             add_info.update(kv_xfer_params)
 
-                    return kv_xfer_params, None
+                    return kv_xfer_params, ec_xfer_params
 
             # 3. Standard Freeing
             delay_free_blocks |= connector_delay_free_blocks
             if not delay_free_blocks:
                 self._free_blocks(request)
 
-            return kv_xfer_params, None
+            return kv_xfer_params, ec_xfer_params
         finally:
             self._free_input_coordinator_request(request_id)
             # Normal completion runs through here, not finish_requests()

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -312,6 +312,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             ec_meta = self.ec_connector.build_connector_meta(scheduler_output)
             scheduler_output.ec_connector_metadata = ec_meta
 
+        # Advance the fence only for non-empty steps (those that actually
+        # write KV and have their output processed later in
+        # update_from_output). Must precede _update_after_schedule, which
+        # stamps request.last_sched_seq from the advanced value; the other
+        # half drains in update_from_output. The fallback path above gets
+        # both halves from super().schedule(). getattr: __new__-constructed
+        # test schedulers carry no defer_block_free attribute.
+        if getattr(self, "defer_block_free", False) and total_num_scheduled_tokens > 0:
+            self.sched_step_seq += 1
+
         # Update internal state (advance num_computed_tokens, free encoder inputs,
         # etc.)
         self._update_after_schedule(scheduler_output)
@@ -348,8 +358,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
+        ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
 
         cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
+
+        # Every GPU write enqueued by this and earlier steps has completed, so
+        # it is safe to return deferred-free blocks to the pool. This is the
+        # update-side half of the upstream v0.28 deferred-free fence; the schedule
+        # half advances sched_step_seq in the fast path above.
+        # getattr: __new__-constructed test schedulers carry no
+        # defer_block_free attribute.
+        if (
+            getattr(self, "defer_block_free", False)
+            # getattr again: SchedulerOutput is mocked field-by-field in the
+            # scheduler unit tests, and dataclass fields without defaults are
+            # invisible to MagicMock(spec=...).
+            and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0
+        ):
+            self.processed_step_seq += 1
+            self._drain_deferred_frees()
+
         perf_stats: PerfStats | None = None
         if self.perf_metrics and self.perf_metrics.is_enabled():
             perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
@@ -441,6 +469,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             new_logprobs = None
             new_token_ids = generated_token_ids
             kv_transfer_params = None
+            ec_transfer_params = None
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             mm_output = mm_outputs[req_index] if mm_outputs else None
             status_before_stop = request.status
@@ -477,6 +506,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # (does not affect protocol)
                 stopped = True
 
+            # Finalize prefill stats BEFORE stop handling (upstream v0.28
+            # order): _free_request below releases the KV blocks, after which
+            # estimate_cached_tokens(request) reports 0. kv_transfer_params is
+            # omitted from the emission predicate here because it only becomes
+            # non-None when stopped is already True.
+            prefill_stats = None
+            if new_token_ids or mm_output is not None or pooler_output is not None or stopped:
+                prefill_stats = request.take_prefill_stats()
+                if prefill_stats is not None:
+                    prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request))
+
             if stopped:
                 if model_runner_output.routed_experts is not None:
                     routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
@@ -488,7 +528,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     if self.chunk_transfer_adapter:
                         self.chunk_transfer_adapter.segment_finished_requests.discard(req_id)
                 if finished:
-                    kv_transfer_params, _ = self._free_request(request)
+                    kv_transfer_params, ec_transfer_params = self._free_request(request)
                     if self.chunk_transfer_adapter is not None:
                         self.chunk_transfer_adapter.cleanup(
                             request.request_id,
@@ -530,8 +570,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
-                    prefill_stats=request.take_prefill_stats(),
+                    prefill_stats=prefill_stats,
                     kv_transfer_params=kv_transfer_params,
+                    ec_transfer_params=ec_transfer_params,
                     routed_experts=routed_experts,
                     num_nans_in_logits=request.num_nans_in_logits,
                     is_segment_finished=is_segment_finished,
@@ -550,8 +591,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             finished = self._handle_stopped_request(request)
             is_segment_finished = not finished
             kv_transfer_params = None
+            ec_transfer_params = None
             if finished:
-                kv_transfer_params, _ = self._free_request(request)
+                kv_transfer_params, ec_transfer_params = self._free_request(request)
                 if self.chunk_transfer_adapter is not None:
                     self.chunk_transfer_adapter.cleanup(
                         request.request_id,
@@ -565,6 +607,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 finish_reason=finish_reason,
                 stop_reason=request.stop_reason,
                 kv_transfer_params=kv_transfer_params,
+                ec_transfer_params=ec_transfer_params,
                 is_segment_finished=is_segment_finished,
             )
             stopped_running_reqs.add(request)
@@ -590,6 +633,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
+        # EC Connector: update state from worker-side EC connector output.
+        # Use getattr for safety with test __new__/SimpleNamespace code paths.
+        if getattr(self, "ec_connector", None) is not None and ec_connector_output:
+            self.ec_connector.update_connector_output(ec_connector_output)
+
         kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
         self._publish_kv_cache_events()
 

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -472,6 +472,7 @@ class OmniSchedulerMixin:
         stop_reason: Any = None,
         prefill_stats: Any = None,
         kv_transfer_params: Any = None,
+        ec_transfer_params: Any = None,
         routed_experts: Any = None,
         num_nans_in_logits: int = 0,
         is_segment_finished: bool | None = False,
@@ -490,6 +491,7 @@ class OmniSchedulerMixin:
             events=request.take_events(),
             prefill_stats=prefill_stats,
             kv_transfer_params=kv_transfer_params,
+            ec_transfer_params=ec_transfer_params,
             trace_headers=request.trace_headers,
             routed_experts=routed_experts,
             num_nans_in_logits=num_nans_in_logits,

vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import torch
@@ -91,6 +94,21 @@ class MoTQKVParallelLinear(QKVParallelLinear):
         # boilerplate Module subclass just for the VAE pathway.
         self.gen_exp = torch.nn.Module()
 
+        # ``gen_exp`` participates in the same tensor-parallel layout as this
+        # layer.  vLLM's online quantizers inspect these attributes while
+        # processing weights (for example, FP8 uses them to share scales across
+        # TP ranks), so the lightweight container must expose the parallel
+        # linear metadata rather than only holding parameters.
+        for attr in (
+            "tp_rank",
+            "tp_size",
+            "input_size",
+            "input_size_per_partition",
+            "output_size",
+            "output_size_per_partition",
+        ):
+            setattr(self.gen_exp, attr, getattr(self, attr))
+
         # Use the same weight_loader as text
         vae_weight_loader = (
             self.weight_loader_v2

vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import torch
@@ -84,6 +87,21 @@ class MoTRowParallelLinear(RowParallelLinear):
         # boilerplate Module subclass just for the VAE pathway.
         self.gen_exp = torch.nn.Module()
 
+        # ``gen_exp`` participates in the same tensor-parallel layout as this
+        # layer.  vLLM's online quantizers inspect these attributes while
+        # processing weights (for example, FP8 uses them to share scales across
+        # TP ranks), so the lightweight container must expose the parallel
+        # linear metadata rather than only holding parameters.
+        for attr in (
+            "tp_rank",
+            "tp_size",
+            "input_size",
+            "input_size_per_partition",
+            "output_size",
+            "output_size_per_partition",
+        ):
+            setattr(self.gen_exp, attr, getattr(self, attr))
+
         # Select weight_loader consistent with text
         vae_weight_loader = (
             self.weight_loader_v2

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -25,7 +25,12 @@ from vllm.logger import init_logger
 from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
 from vllm.utils.system_utils import get_mp_context
 from vllm.v1.engine.core import EngineCoreProc
-from vllm.v1.engine.utils import CoreEngine, EngineZmqAddresses, wait_for_engine_startup
+from vllm.v1.engine.utils import (
+    CoreEngine,
+    CoreEngineLaunch,
+    EngineZmqAddresses,
+    wait_for_engine_startup,
+)
 from vllm.v1.utils import shutdown
 
 from vllm_omni.diffusion.data import DiffusionRequestAbortedError
@@ -764,7 +769,6 @@ class StageDiffusionProcManager:
             with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
                 wait_for_engine_startup(
                     handshake_socket,
-                    self.addresses,
                     [CoreEngine(index=0, local=True)],
                     SimpleNamespace(
                         data_parallel_size_local=1,
@@ -773,8 +777,17 @@ class StageDiffusionProcManager:
                     ),
                     False,
                     None,
-                    self,
-                    None,
+                    CoreEngineLaunch(
+                        engine_manager=self,
+                        coordinator=None,
+                        addresses=self.addresses,
+                        tensor_queue=None,
+                        # StageDiffusionProcManager is not a CoreEngineProcManager,
+                        # so upstream's isinstance-gated sentinel registration would
+                        # be skipped; watch the subprocess directly so a proc death
+                        # during handshake is still detected.
+                        watched_frontend_processes=[self.proc],
+                    ),
                 )
         except Exception:
             shutdown([self.proc])

vllm_omni/engine/stage_engine_startup.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Helpers for launching and handshaking omni engine cores."""
 
 from __future__ import annotations
@@ -22,6 +25,7 @@ from vllm.utils.network_utils import get_open_ports_list, zmq_socket_ctx
 from vllm.v1.engine.coordinator import DPCoordinator
 from vllm.v1.engine.utils import (
     CoreEngine,
+    CoreEngineLaunch,
     CoreEngineProcManager,
     EngineZmqAddresses,
     get_engine_zmq_addresses,
@@ -833,13 +837,16 @@ def connect_remote_engine_cores(
             )
             wait_for_engine_startup(
                 handshake_socket,
-                addresses,
                 engines_to_handshake,
                 vllm_config.parallel_config,
                 False,  # coordinated_dp
                 vllm_config.cache_config,
-                None,  # proc_manager (remote — no local procs)
-                None,  # coord_process
+                CoreEngineLaunch(
+                    engine_manager=None,  # remote — no local procs
+                    coordinator=None,
+                    addresses=addresses,
+                    tensor_queue=None,
+                ),
             )
     finally:
         omni_master_server.release_route_port_reservations(
@@ -875,13 +882,16 @@ def connect_remote_diffusion_proc(
             yield StageReplicaResources(addresses=addresses)
             wait_for_engine_startup(
                 handshake_socket,
-                addresses,
                 [CoreEngine(index=0, local=False)],
                 _single_diffusion_parallel_config(local_client=False),
                 False,
                 None,
-                None,
-                None,
+                CoreEngineLaunch(
+                    engine_manager=None,  # remote — no local procs
+                    coordinator=None,
+                    addresses=addresses,
+                    tensor_queue=None,
+                ),
             )
     finally:
         omni_master_server.release_route_port_reservations(
@@ -1077,13 +1087,16 @@ def _launch_omni_core_engines(
             yield local_engine_manager, coordinator, addresses
             wait_for_engine_startup(
                 handshake_socket,
-                addresses,
                 engines_to_handshake,
                 parallel_config,
                 parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe,
                 vllm_config.cache_config,
-                local_engine_manager,
-                coordinator.proc if coordinator else None,
+                CoreEngineLaunch(
+                    engine_manager=local_engine_manager,
+                    coordinator=coordinator,
+                    addresses=addresses,
+                    tensor_queue=None,
+                ),
             )
     finally:
         omni_master_server.release_route_port_reservations(
@@ -1166,13 +1179,16 @@ def launch_stage_replica(
         )
         wait_for_engine_startup(
             handshake_socket,
-            addresses,
             engines_to_handshake,
             vllm_config.parallel_config,
             False,  # coordinated_dp
             vllm_config.cache_config,
-            engine_manager,
-            None,  # coordinator_proc
+            CoreEngineLaunch(
+                engine_manager=engine_manager,
+                coordinator=None,
+                addresses=addresses,
+                tensor_queue=None,
+            ),
         )
 
 

vllm_omni/entrypoints/openai/api_server.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import asyncio
 import base64
 import dataclasses
@@ -64,6 +64,7 @@ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAISer
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
+from vllm.entrypoints.serve import create_error_response
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
@@ -75,7 +76,6 @@ from vllm.entrypoints.serve.utils.api_utils import (
     validate_json_request,
     with_cancellation,
 )
-from vllm.entrypoints.serve.utils.error_response import create_error_response
 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
 from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config
@@ -997,11 +997,9 @@ async def omni_init_app_state(
     )
 
     # Warm up chat template processing to avoid first-request latency
-    # Upstream f5ffc59b6a moved warmup onto OnlineRenderer. Accelerator
-    # images can temporarily lag that renderer API, where warmup is optional.
-    renderer_warmup = getattr(state.online_renderer, "warmup", None)
-    if renderer_warmup is not None:
-        renderer_warmup()
+    # Upstream f5ffc59b6a removed OpenAIServingChat.warmup() and moved the
+    # warmup onto the renderer (OnlineRenderer.warmup()); mirror upstream.
+    state.online_renderer.warmup()
 
     state.openai_serving_completion = (
         OpenAIServingCompletion(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import base64
 import json
@@ -19,9 +22,11 @@ from vllm.entrypoints.chat_utils import (
     ChatCompletionMessageParam,
     ChatTemplateContentFormatOption,
     ConversationMessage,
-    get_history_tool_calls_cnt,
     make_tool_call_id,
 )
+from vllm.parser.utils import (
+    count_chat_history_tool_calls as get_history_tool_calls_cnt,
+)
 
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -86,14 +91,7 @@ from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.multimodal.media.connector import MediaConnector
 from vllm.outputs import RequestOutput
-
-try:
-    from vllm.parser.mistral import MistralToolCall
-except ImportError:
-    # Accelerator images can lag the core vLLM release by one version.
-    from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
-
-
+from vllm.parser.mistral import MistralToolCall
 from vllm.reasoning import ReasoningParser
 from vllm.renderers import BaseRenderer, merge_kwargs
 from vllm.renderers.inputs import TokPrompt

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import glob
 import os
 from collections.abc import Iterable
@@ -19,9 +22,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
-from vllm.model_executor.models.vision import (
-    get_llm_pos_ids_for_vision,
-)
 
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -33,6 +33,9 @@ from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+from vllm_omni.model_executor.layers.rotary_embedding.mrope import (
+    OmniMRotaryEmbedding,
+)
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
@@ -43,6 +46,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.platforms import current_omni_platform
 
+# `get_llm_pos_ids_for_vision` was removed from upstream vLLM (commit
+# 0cf49a5d15); vllm-omni keeps an identical copy as a static method on
+# OmniMRotaryEmbedding. Bind it at module level to keep call sites unchanged.
+get_llm_pos_ids_for_vision = OmniMRotaryEmbedding._get_llm_pos_ids_for_vision
+
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
 

vllm_omni/outputs/output_processor.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import fields as dataclass_fields
 from typing import Any
 
@@ -148,6 +151,8 @@ class OmniRequestState(RequestState):
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
         kv_transfer_params: dict[str, Any] | None = None,
+        ec_transfer_params: dict[str, Any] | None = None,
+        *,
         routed_experts: Any = None,
     ) -> OmniRequestOutput | PoolingRequestOutput | None:
         """Create a request output from generation results.
@@ -162,8 +167,13 @@ class OmniRequestState(RequestState):
             finish_reason: Optional finish reason indicating why generation stopped
             stop_reason: Optional stop reason (token ID or stop string)
             kv_transfer_params: Optional KV cache transfer parameters
+            ec_transfer_params: Optional encoder-cache transfer parameters
+                (6th positional, matching upstream RequestState so that
+                super().process_outputs() calls do not misroute it).
             routed_experts: Optional MoE routed-expert ids for this step,
                 attached to the completion output for generation stages
+                (omni-specific keyword; upstream moved this accumulation
+                into RequestState.routed_experts_chunks).
 
         Returns:
             OmniRequestOutput or PoolingRequestOutput if output should be
@@ -177,6 +187,7 @@ class OmniRequestState(RequestState):
                 finish_reason,
                 stop_reason,
                 kv_transfer_params,
+                ec_transfer_params,
             )
 
         is_delta = self.output_kind == RequestOutputKind.DELTA
@@ -230,6 +241,7 @@ class OmniRequestState(RequestState):
             outputs,
             finished,
             kv_transfer_params,
+            ec_transfer_params,
         )
 
     def _new_completion_output(
@@ -309,6 +321,7 @@ class OmniRequestState(RequestState):
         outputs: list,
         finished: bool,
         kv_transfer_params: dict[str, Any] | None = None,
+        ec_transfer_params: dict[str, Any] | None = None,
     ) -> RequestOutput | PoolingRequestOutput:
         """Create request output, handling no-detokenizer generation stages.
 
@@ -323,6 +336,7 @@ class OmniRequestState(RequestState):
                 outputs,
                 finished,
                 kv_transfer_params,
+                ec_transfer_params,
             )
 
         # No-detokenizer path: build RequestOutput directly.
@@ -341,7 +355,9 @@ class OmniRequestState(RequestState):
             outputs=outputs,
             finished=finished,
             kv_transfer_params=kv_transfer_params,
+            ec_transfer_params=ec_transfer_params,
             num_cached_tokens=self.num_cached_tokens,
+            num_cache_creation_tokens=self.num_cache_creation_tokens,
             metrics=self.stats,
         )
 
@@ -560,6 +576,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             finish_reason = eco.finish_reason
             stop_reason = eco.stop_reason
             kv_transfer_params = eco.kv_transfer_params
+            ec_transfer_params = getattr(eco, "ec_transfer_params", None)
             routed_experts = eco.routed_experts
             self._update_stats_from_output(
                 req_state,
@@ -570,6 +587,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             prefill_stats = getattr(eco, "prefill_stats", None)
             if prefill_stats is not None:
                 req_state.num_cached_tokens = prefill_stats.num_cached_tokens
+                req_state.num_cache_creation_tokens = prefill_stats.num_cache_creation_tokens
             req_state.is_prefilling = False
 
             is_non_final_audio_chunk = (
@@ -583,7 +601,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 finish_reason,
                 stop_reason,
                 kv_transfer_params,
-                routed_experts,
+                ec_transfer_params,
+                routed_experts=routed_experts,
             ):
                 if req_state.queue is not None:
                     req_state.queue.put(request_output)

vllm_omni/quantization/int8_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """INT8 quantization config for diffusion transformers.
 
 Supports both online (dynamic) and offline (checkpoint) INT8 quantization
@@ -26,10 +26,12 @@ 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.utils.quant_utils import (
     is_layer_skipped,
 )
+from vllm.model_executor.model_loader.reload.meta import (
+    CopyCounter as CopyNumelCounter,
+)
 from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
 from vllm.model_executor.parameter import (
     ChannelQuantScaleParameter,

vllm_omni/quantization/mxfp8_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """W8A8 MXFP8 (Microscaling FP8) online/offline quantization for diffusion transformers.
 
 Architecture:
@@ -49,8 +49,10 @@ 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.utils.quant_utils import is_layer_skipped
+from vllm.model_executor.model_loader.reload.meta import (
+    CopyCounter as CopyNumelCounter,
+)
 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

vllm_omni/quantization/svdquant_config.py

@@ -157,7 +157,10 @@ class DiffusionSVDQuantConfig(QuantizationConfig):
             prefix,
             self.modules_to_not_convert,
             self.packed_modules_mapping,
-            skip_with_substr=True,
+            # vLLM 0.28: skip_with_substr was replaced by match_mode
+            # ("substring" preserves the v0.27 skip_with_substr=True
+            # behavior verbatim).
+            match_mode="substring",
         ):
             return UnquantizedLinearMethod()
         return DiffusionSVDQuantLinearMethod(self)

vllm_omni/request.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
@@ -103,6 +106,7 @@ class OmniRequest(Request):
             additional_information=request.additional_information,
             model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None),
             resumable=request.resumable,
+            session_id=request.session_id,
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """AR GPU Model Runner for vLLM-Omni.
 
 Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
@@ -169,6 +172,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         async_output_copy_stream = kwargs.pop("async_output_copy_stream")
         vocab_size = kwargs.pop("vocab_size")
         routed_experts = kwargs.pop("routed_experts", None)
+        num_nans = kwargs.pop("num_nans", None)
         # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
         # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
         # feature but must consume the kwarg to prevent TypeError from stray
@@ -186,6 +190,11 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         self._logprobs_tensors = logprobs_tensors
         self._routed_experts = routed_experts
         self._has_fault: torch.Tensor | None = None
+        # Upstream b1e12d142d (PR #51304) added device-side NaN-in-logits
+        # counts (num_nans) to AsyncGPUModelRunnerOutput. Omni keeps the
+        # counts on the async copy stream and lets super().get_output()
+        # populate num_nans_in_logits from the CPU copy.
+        self._num_nans = num_nans
 
         default_stream = torch.cuda.current_stream()
         with torch.cuda.stream(async_output_copy_stream):
@@ -199,6 +208,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
             self._routed_experts_cpu = (
                 self._routed_experts.to_cpu_nonblocking() if self._routed_experts is not None else None
             )
+            self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
             self.async_copy_ready_event.record()
 
         self._model_runner_output_builder = model_runner_output_builder
@@ -242,6 +252,12 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         # (e.g. unit tests using object.__new__).
         if not hasattr(self, "_has_fault"):
             self._has_fault = None
+        # Upstream b1e12d142d (PR #51304) also touches _num_nans/_num_nans_cpu
+        # in get_output(). Guard them the same way for object.__new__ tests.
+        if not hasattr(self, "_num_nans"):
+            self._num_nans = None
+        if not hasattr(self, "_num_nans_cpu"):
+            self._num_nans_cpu = None
         with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
             return super().get_output()
 
@@ -1030,9 +1046,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
             num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
 
-            logits_indices, spec_decode_metadata = self._prepare_inputs(
-                scheduler_output,
-                num_scheduled_tokens_np,
+            logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs(
+                scheduler_output, num_scheduled_tokens_np
             )
 
             cascade_attn_prefix_lens = None
@@ -1126,6 +1141,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 max_query_len=max_num_scheduled_tokens,
                 ubatch_slices=ubatch_slices_attn,
                 logits_indices=logits_indices,
+                max_num_sampled_tokens=max_num_sampled_tokens,
                 use_spec_decode=use_spec_decode,
                 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
@@ -1165,15 +1181,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
             )
 
-        # Set cudagraph mode to none if calc_kv_scales is true.
-        # KV scales calculation involves dynamic operations that are incompatible
-        # with CUDA graph capture.
-        if self.calculate_kv_scales:
-            cudagraph_mode = CUDAGraphMode.NONE
-            runner_assisted_full_attn = False
-            # Mark KV scales as calculated after the first forward pass
-            self.calculate_kv_scales = False
-
         runner_assisted_context_enabled = False
         if runner_assisted_full_attn:
             runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context(
@@ -2101,6 +2108,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
             (
                 num_nans_in_logits,
+                num_nans,
                 logprobs_lists,
                 valid_sampled_token_ids,
                 prompt_logprobs_dict,
@@ -2218,6 +2226,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 invalid_req_indices=invalid_req_indices,
                 async_output_copy_stream=self.async_output_copy_stream,
                 vocab_size=self.input_batch.vocab_size,
+                num_nans=num_nans,
             )
             if use_async_omni_output:
                 async_output = async_output_cls(

vllm_omni/worker/gpu_generation_model_runner.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Code2Wav GPU Model Runner for vLLM-Omni.
 
 Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
@@ -192,9 +195,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
             num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
 
-            logits_indices, spec_decode_metadata = self._prepare_inputs(
-                scheduler_output,
-                num_scheduled_tokens_np,
+            logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs(
+                scheduler_output, num_scheduled_tokens_np
             )
 
             cascade_attn_prefix_lens = None
@@ -278,6 +280,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 max_query_len=max_num_scheduled_tokens,
                 ubatch_slices=ubatch_slices_attn,
                 logits_indices=logits_indices,
+                max_num_sampled_tokens=max_num_sampled_tokens,
                 use_spec_decode=use_spec_decode,
                 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
@@ -309,14 +312,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if getattr(self.model, "requires_request_ids", False):
                 model_kwargs["request_ids"] = list(req_ids)
 
-        # Set cudagraph mode to none if calc_kv_scales is true.
-        # KV scales calculation involves dynamic operations that are incompatible
-        # with CUDA graph capture.
-        if self.calculate_kv_scales:
-            cudagraph_mode = CUDAGraphMode.NONE
-            # Mark KV scales as calculated after the first forward pass
-            self.calculate_kv_scales = False
-
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         # When spec decode is enabled, defer connector finalization