From 651b03225b69fe825525e1cfa00300fe856fcea0 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Thu, 27 Aug 2026 14:12:55 +0800 Subject: [PATCH] [Rebase] Rebase to vllm 0.28.0 (#6606) Signed-off-by: tzhouam Co-authored-by: Claude Opus 5 (1M context) --- docker/Dockerfile.ci | 13 +- docker/Dockerfile.cuda | 2 +- docker/Dockerfile.rocm | 4 +- docs/getting_started/installation/README.md | 2 +- docs/getting_started/installation/gpu.md | 2 +- .../installation/gpu/cuda.inc.md | 18 +- .../installation/gpu/rocm.inc.md | 14 +- docs/getting_started/quickstart.md | 6 +- examples/__init__.py | 9 + .../test_omni_ar_scheduler_stale_drain.py | 1 + .../sched/test_omni_ar_scheduler_streaming.py | 2 + ...mni_generation_scheduler_update_session.py | 1 + .../test_omni_sched_deferred_free_fence.py | 168 ++++++++++++++++ .../test_omni_sched_ec_request_finish.py | 180 ++++++++++++++++++ .../test_omni_sched_prefill_stats_finalize.py | 126 ++++++++++++ .../diffusion/kernels/mot/test_mot_linear.py | 68 +++++++ .../test_lingbot_world_camera.py | 4 +- .../test_chunk_transfer_adapter.py | 3 + .../e2e/offline_inference/test_flux2_klein.py | 2 + .../test_async_omni_engine_stage_init.py | 4 + tests/engine/test_output_processor.py | 51 ++++- tests/engine/test_single_stage_mode.py | 5 +- .../test_omni_snapshot_download.py | 14 +- tests/worker/test_gpu_ar_model_runner.py | 4 + vllm_omni/config/endpoint_policy.py | 4 +- vllm_omni/core/sched/omni_ar_scheduler.py | 58 +++++- .../core/sched/omni_generation_scheduler.py | 54 +++++- vllm_omni/core/sched/omni_scheduler_mixin.py | 2 + .../layers/mot/mot_qkv_parallel_linear.py | 18 ++ .../layers/mot/mot_row_parallel_linear.py | 18 ++ vllm_omni/diffusion/stage_diffusion_proc.py | 21 +- vllm_omni/engine/stage_engine_startup.py | 40 ++-- vllm_omni/entrypoints/openai/api_server.py | 12 +- vllm_omni/entrypoints/openai/serving_chat.py | 16 +- .../models/qwen2_5_omni/qwen2_5_omni.py | 14 +- vllm_omni/outputs/output_processor.py | 21 +- vllm_omni/quantization/int8_config.py | 6 +- vllm_omni/quantization/mxfp8_config.py | 6 +- vllm_omni/quantization/svdquant_config.py | 5 +- vllm_omni/request.py | 4 + vllm_omni/worker/gpu_ar_model_runner.py | 33 ++-- .../worker/gpu_generation_model_runner.py | 17 +- 42 files changed, 931 insertions(+), 121 deletions(-) create mode 100644 examples/__init__.py create mode 100644 tests/core/sched/test_omni_sched_deferred_free_fence.py create mode 100644 tests/core/sched/test_omni_sched_ec_request_finish.py create mode 100644 tests/core/sched/test_omni_sched_prefill_stats_finalize.py diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index a7e6ccad0..b84b344a8 100644 --- a/docker/Dockerfile.ci +++ b/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 diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index f2e80eaf3..3881ca2e5 100644 --- a/docker/Dockerfile.cuda +++ b/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 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 4283efc77..597f5d134 100644 --- a/docker/Dockerfile.rocm +++ b/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}} diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md index b49cee794..448ca81f2 100644 --- a/docs/getting_started/installation/README.md +++ b/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: diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 59d2aa482..196b6a638 100644 --- a/docs/getting_started/installation/gpu.md +++ b/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" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index a68bfdb58..a8f795861 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/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 . ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index c44407299..e0f1614ae 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/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] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index a5a02f5a7..439c7aedc 100644 --- a/docs/getting_started/quickstart.md +++ b/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 diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 000000000..78a85387e --- /dev/null +++ b/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``. +""" diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index fb92414ef..7c72602cf 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/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 diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index bf35bfdbc..cd5a11672 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/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 diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 247f0ce84..3e1b53afe 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/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 diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py new file mode 100644 index 000000000..376c64cb3 --- /dev/null +++ b/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 diff --git a/tests/core/sched/test_omni_sched_ec_request_finish.py b/tests/core/sched/test_omni_sched_ec_request_finish.py new file mode 100644 index 000000000..81b6577fd --- /dev/null +++ b/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) diff --git a/tests/core/sched/test_omni_sched_prefill_stats_finalize.py b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py new file mode 100644 index 000000000..d436367db --- /dev/null +++ b/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) diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index babbc7abc..49622aecd 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/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) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py index 8b30237c4..121364c07 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py +++ b/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, ) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index a5537bdae..ca518bd79 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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) diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py index 845cf2789..ed0cbaf91 100644 --- a/tests/e2e/offline_inference/test_flux2_klein.py +++ b/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() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 1b3c3be40..6203ed726 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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] = {} diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 23ac601b3..076e8ad91 100644 --- a/tests/engine/test_output_processor.py +++ b/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 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 039b5119a..b576e72d4 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/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 diff --git a/tests/entrypoints/test_omni_snapshot_download.py b/tests/entrypoints/test_omni_snapshot_download.py index 58dd1d19d..eae774178 100644 --- a/tests/entrypoints/test_omni_snapshot_download.py +++ b/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) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 539329d24..5e3cfe865 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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"], diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index d22c0acc1..06731603e 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/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): diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index abecfdeef..02be8c581 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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() diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index d572453fd..533f48062 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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() diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 4b7c89539..f9be46248 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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, diff --git a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py index ce83d8842..235fc3342 100644 --- a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py +++ b/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 diff --git a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py index 2811b110f..6e846af47 100644 --- a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 224a88199..4419e51dc 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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]) diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 5dacc31df..e9622b83a 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/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, + ), ) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 57adaad08..0adc9468e 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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( diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 48812324f..b7f99e3b3 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index ce6617695..2f653b12f 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/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 diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 9558d66e5..c6ad6ca28 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/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) diff --git a/vllm_omni/quantization/int8_config.py b/vllm_omni/quantization/int8_config.py index 01e57e708..de31fbc3a 100644 --- a/vllm_omni/quantization/int8_config.py +++ b/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, diff --git a/vllm_omni/quantization/mxfp8_config.py b/vllm_omni/quantization/mxfp8_config.py index 433531d56..69e51dc6b 100644 --- a/vllm_omni/quantization/mxfp8_config.py +++ b/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 diff --git a/vllm_omni/quantization/svdquant_config.py b/vllm_omni/quantization/svdquant_config.py index 80478361b..ad22592af 100644 --- a/vllm_omni/quantization/svdquant_config.py +++ b/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) diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 0606bcfc8..a6da16696 100644 --- a/vllm_omni/request.py +++ b/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, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 029254d6c..cd03cf3fc 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 73fe406e6..a1e136e53 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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