diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 71ed75119..993854541 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -102,6 +102,7 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback" CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN" CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi' +CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py" # Bootstrap Buildkite ``if`` expressions. # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN). @@ -620,11 +621,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()): commands = merged.get("commands") if isinstance(commands, list): - merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands] + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands] elif commands is None: - merged["commands"] = [CUDA_HF_TOKEN_EXPORT] + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK] else: - merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands] + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands] # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out. if "retry" in step: merged["retry"] = step["retry"] diff --git a/.buildkite/cuda/scripts/check_vllm_runtime.py b/.buildkite/cuda/scripts/check_vllm_runtime.py new file mode 100644 index 000000000..813d9f6d8 --- /dev/null +++ b/.buildkite/cuda/scripts/check_vllm_runtime.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker.""" + +from __future__ import annotations + +import importlib +import re +from pathlib import Path + + +def main() -> None: + dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci" + match = re.search( + r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$", + dockerfile.read_text(), + re.MULTILINE, + ) + if match is None: + raise RuntimeError(f"Missing vLLM release version in {dockerfile}") + expected_version = match.group(1) + + import torch + import vllm + + assert vllm.__version__ == expected_version, vllm.__version__ + assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker" + for module in ("flashinfer", "vllm._C_stable_libtorch"): + importlib.import_module(module) + print( + f"Verified vLLM CUDA runtime: vllm={vllm.__version__} " + f"expected={expected_version} torch={torch.__version__} " + f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}" + ) + + +if __name__ == "__main__": + main() diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 376e322ed..86e794c7a 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -48,6 +48,8 @@ steps: - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" source_file_dependencies: omni_qwen3_omni_doc timeout_in_minutes: 90 + artifact_paths: + - audio_*.wav commands: - pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model" @@ -156,6 +158,11 @@ steps: - label: ":full_moon: TTS · Function Test with L4" source_file_dependencies: tts_qwen3_tts_function timeout_in_minutes: 120 + env: + VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: . + artifact_paths: + - failed_speech_*.wav + - failed_speech_*.json commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_USE_DEEP_GEMM="0" diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index cb62fd310..949fa0970 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,10 +1,7 @@ -# Target vLLM version: v0.30.0 (tag @ ced6857afa0e). -# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer -# runtime. Do not reinstall a commit wheel or repair dependencies on top of -# an older base image at a release boundary. For a future unreleased target, -# restore an explicit wheel pin and validate its complete runtime stack. +# Use the official release image with its matching CUDA/PyTorch stack. +ARG VLLM_VERSION=0.31.0 ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.30.0 +ARG VLLM_BASE_TAG=v${VLLM_VERSION} # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents @@ -12,6 +9,7 @@ ARG VLLM_BASE_TAG=v0.30.0 # everything except the actual project dir, since that should not # change much across branches. FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps +ARG VLLM_VERSION # Keep APP_DIR outside /workspace so k8s Buildkite workspace mounts cannot # shadow the editable install target (#5364). Docker WORKDIR creates /opt if # needed. @@ -36,12 +34,17 @@ RUN apt-get update && \ # for CI COPY pyproject.toml setup.py ./ COPY requirements/ requirements/ -ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 +# Scope the override to Omni; dependency source builds need their own versions. +ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0 # Editable install: deps + project metadata are cached in this layer, while # the package itself is linked to APP_DIR (populated later by COPY --link). # Do NOT use a non-editable install here — that copies an empty stub into # site-packages and breaks ``import vllm_omni.entrypoints``. RUN uv pip install --system -e ".[dev]" +# Fail the image build if dependency resolution replaced the selected vLLM. +# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start +# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py. +RUN python3 -c "import vllm; assert vllm.__version__ == '${VLLM_VERSION}', vllm.__version__" # The final stage just copies the working dir on top of the # dependencies image. This uses --link so that we can skip diff --git a/docker/Dockerfile.ci.dockerignore b/docker/Dockerfile.ci.dockerignore index ab137c2b5..4564eb75e 100644 --- a/docker/Dockerfile.ci.dockerignore +++ b/docker/Dockerfile.ci.dockerignore @@ -1,6 +1,6 @@ # NOTE: this is largely taken from vLLM -# Exclude .git for CI builds — version is set via SETUPTOOLS_SCM_PRETEND_VERSION +# Exclude .git for CI builds — version is set via VLLM_OMNI_VERSION_OVERRIDE .git .venv /build diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index cba47dc25..69b598ab2 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.30.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.31.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 1639bcfeb..379262afa 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,7 +1,7 @@ # Keep the ROCm runtime aligned with the vLLM release targeted by # docker/Dockerfile.ci. A stale base image can build successfully while failing # at runtime when vLLM-Omni imports APIs added by the newer vLLM release. -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 # Keep the dependency-heavy portion independent from the source tree. The # registry cache is keyed on these inputs, while the final source copy can be @@ -12,11 +12,11 @@ FROM ${BASE_IMAGE} AS deps # BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0. # VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when # USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's -# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An +# release version. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An # intentional custom source pin must update that regression in the same change, # and should return to the stable release as soon as one is available. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0 +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.31.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} @@ -75,7 +75,8 @@ WORKDIR ${COMMON_WORKDIR}/vllm-omni RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py COPY pyproject.toml setup.py ./ COPY requirements/ requirements/ -ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 +# Scope the override to Omni; dependency source builds need their own versions. +ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \ --no-cache-dir -e ".[dev]" --no-build-isolation && \ diff --git a/docs/configuration/README.md b/docs/configuration/README.md index fc12c2663..b8640ecad 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -2,7 +2,7 @@ This section lists the most common options for running vLLM-Omni. -For options within a vLLM Engine, please refer to the [vLLM 0.30 configuration guide](https://docs.vllm.ai/en/v0.30.0/configuration/index.html). +For options within a vLLM Engine, please refer to the [vLLM 0.31 configuration guide](https://docs.vllm.ai/en/v0.31.0/configuration/index.html). Each model defines fixed topology in a registered `PipelineConfig` and runtime overrides in a deploy YAML. diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 9cc0abcad..5a26c00a3 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -186,7 +186,7 @@ their keys only. ## Inherited vLLM variables vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to -the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.0/configuration/env_vars.html) +the [vLLM 0.31 environment-variable reference](https://docs.vllm.ai/en/v0.31.0/configuration/env_vars.html) for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm, XPU, ModelScope, and FlashInfer workspace settings. diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md index 507a0bd96..e82075aba 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. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0. + vLLM-Omni is released against the matching upstream vLLM major/minor version. This source checkout targets vLLM 0.31.0. The published vLLM-Omni 0.30.0 wheel instructions use vLLM 0.30.0; choose the installation path that matches your Omni revision. vLLM-Omni supports the following hardware platforms: diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index d3704b0b1..fb7db8717 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -1,11 +1,13 @@ # --8<-- [start:requirements] -- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.) +- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, etc.) # --8<-- [end:requirements] # --8<-- [start:set-up-using-python] -vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x. +vLLM-Omni depends on the matching major/minor release of vLLM. This source +checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install +the published 0.30.0 release with its matching vLLM dependency. !!! note PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See for more details. @@ -41,13 +43,17 @@ uv pip install 'vllm-omni[demo]==0.30.0' # --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.30.0 release version of the library +If you do not need to modify source code of vLLM, install the final 0.31.0 +release wheel from its commit-pinned index: ```bash -uv pip install vllm==0.30.0 --torch-backend=auto +uv pip install vllm==0.31.0 --torch-backend=auto \ + --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6 ``` -The 0.30.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 default 0.31.0 wheel contains CUDA 12.9 binaries. 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 @@ -68,16 +74,16 @@ 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.30.0 +git checkout v0.31.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.30.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.30.0/vllm-0.30.0-cp38-abi3-manylinux_2_28_x86_64.whl +# Default CUDA 12.9 wheel for v0.31.0 on x86_64 +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash -uv pip install --editable . +VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto ``` Install vllm with command below (If you already have PyTorch). ```bash @@ -128,12 +134,13 @@ docker run --runtime nvidia --gpus 2 \ DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda . ``` -If you want to specify the base vLLM version: +The source build defaults to `vllm/vllm-openai:v0.31.0`. To select that +upstream release image explicitly: ```bash DOCKER_BUILDKIT=1 docker build \ -f docker/Dockerfile.cuda \ - --build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \ + --build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.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 63675ea94..db63f9488 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -10,8 +10,9 @@ For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images. -vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni -0.30.x release line uses vLLM 0.30.x. +vLLM-Omni depends on the matching major/minor release of vLLM. This source +checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install +the published 0.30.0 release with its matching vLLM dependency. The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The `vllm-omni` package does not install vLLM as a dependency, so non-Docker source @@ -49,13 +50,14 @@ uv pip install onnxruntime-rocm #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library +If you do not need to modify source code of vLLM, install the final 0.31.0 +release wheel for ROCm 7.2.3: ```bash -uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723 +uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723 ``` -The pre-built 0.30.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.31.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 @@ -75,7 +77,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.30.0 +git checkout v0.31.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -87,7 +89,7 @@ python3 setup.py develop #### Build docker image The source build defaults to the published upstream base image -`vllm/vllm-openai-rocm:v0.30.0`, aligned with the vLLM release used by CI. +`vllm/vllm-openai-rocm:v0.31.0`, aligned with this source checkout. Older or custom bases must provide the vLLM APIs checked by the Dockerfile's image-build canary. This upstream base is distinct from the prebuilt `vllm/vllm-omni-rocm` images discussed below; their published-tag availability @@ -102,7 +104,7 @@ To select the upstream ROCm base explicitly, pass `BASE_IMAGE`: ```bash DOCKER_BUILDKIT=1 docker build \ -f docker/Dockerfile.rocm \ - --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 \ + --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 \ -t vllm-omni-rocm . ``` diff --git a/docs/getting_started/installation/gpu/xpu.inc.md b/docs/getting_started/installation/gpu/xpu.inc.md index 0232c17bf..49bbbf50f 100644 --- a/docs/getting_started/installation/gpu/xpu.inc.md +++ b/docs/getting_started/installation/gpu/xpu.inc.md @@ -20,12 +20,16 @@ vLLM-Omni currently recommends using the Docker image setup steps below. #### Build docker image ```bash -DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g . +export VLLM_VERSION=v0.31.0 +DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \ + --build-arg VLLM_VERSION="$VLLM_VERSION" . ``` This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:` -base image (v0.30.0 by default), which Docker pulls automatically. To target a different vLLM release, -pass `--build-arg VLLM_VERSION=`. If that tag has not been published yet, build +base image, which Docker pulls automatically. The Dockerfile still defaults to +v0.30.0; the command above explicitly selects v0.31.0 to match this source +checkout. To target a different vLLM release, pass +`--build-arg VLLM_VERSION=`. If that tag has not been published yet, build the base from upstream's own Dockerfile first and point the build at it: ```bash diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index ad29c8447..9bf7613f1 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -71,10 +71,14 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886 # --8<-- [start:installation-main] -You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. +You can also install vLLM-Omni from source if you want to use the latest features +or bug fixes. This source checkout targets vLLM 0.31.0 and needs a matching +vLLM-Ascend runtime. The v0.30.0 images above are for the published Omni v0.30.0 +release. Check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) +for aligned versions and NPU validation before selecting a development image. ```bash -# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main +# Inside a container with matching vLLM and vLLM-Ascend, install Omni from source cd /vllm-workspace git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 0e363c533..e3df99bcd 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,17 +19,20 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.30.0 --torch-backend=auto +uv pip install vllm==0.31.0 --torch-backend=auto \ + --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6 # On ROCm -uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723 +uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni uv pip install -e . ``` -For additional installation methods — please see the [installation guide](installation/README.md). +These commands install vLLM 0.31.0 for this source checkout. The CUDA index is +pinned to the final release commit. For published vLLM-Omni wheels and additional +installation methods, see the [installation guide](installation/README.md). !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. diff --git a/requirements/common.txt b/requirements/common.txt index 9a745cd2c..5a702f66d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -2,7 +2,7 @@ # kernels==0.16.1 requires the LayerRepository version=/revision= integration # introduced in transformers 5.13. Keep the tested upper bound: 5.15 has # a reported model-construction regression; later versions need validation. -# This range also satisfies vLLM 0.30.0's transformers>=5.10.4 requirement. +# This range also satisfies vLLM 0.31.0's transformers>=5.10.4,<5.18.0 requirement. # MOSS-TTS-Nano's list/set weight-loading compatibility is handled by # transformers_keys_to_ignore_compat in model_executor/models/utils.py. transformers >= 5.13.0, < 5.15 diff --git a/tests/buildkite/test_rocm_dockerfile.py b/tests/buildkite/test_rocm_dockerfile.py index 8101b2157..4c40a7c7b 100644 --- a/tests/buildkite/test_rocm_dockerfile.py +++ b/tests/buildkite/test_rocm_dockerfile.py @@ -10,7 +10,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] REPO_ROOT = Path(__file__).resolve().parents[2] AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2" AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh" -CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci" +CUDA_RELEASE_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.cuda" +CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci" ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm" ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore" @@ -30,14 +31,18 @@ def _line_index(lines: list[str], prefix: str) -> int: return matches[0] -def test_rocm_base_tracks_ci_vllm_release() -> None: - ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") +def test_rocm_base_tracks_cuda_vllm_release() -> None: + cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE") + cuda_image, _, cuda_release = cuda_base.rpartition(":") rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE") image_ref, separator, image_tag = rocm_base.rpartition(":") assert separator, f"expected a tagged ROCm base image, got {rocm_base}" - assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref - assert image_tag == ci_release + assert cuda_image == "vllm/vllm-openai" + assert image_ref == "vllm/vllm-openai-rocm" + assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}" + assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image + assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}" def test_rocm_defaults_to_prebuilt_base_image() -> None: @@ -103,11 +108,11 @@ def test_rocm_build_context_excludes_git_history() -> None: assert ".git" in patterns -def test_rocm_source_ref_tracks_ci_vllm_release() -> None: - ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") +def test_rocm_source_ref_tracks_cuda_vllm_release() -> None: + cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION") rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH") - assert rocm_source_ref == ci_release + assert rocm_source_ref == f"v{cuda_release}" def test_rocm_dockerfile_contains_vllm_api_canary() -> None: diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index b39dba2cd..bfdfb85b5 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -172,7 +172,11 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest env_names = {item["name"] for item in container["env"]} assert "VLLM_CI_HF_TOKEN" in env_names assert "HF_TOKEN" not in env_names - assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"] + assert step["commands"] == [ + CUDA_HF_TOKEN_EXPORT, + "python3 .buildkite/cuda/scripts/check_vllm_runtime.py", + "pytest -sv tests/example", + ] def test_mirror_hardwares_l4_preserves_explicit_retry() -> None: @@ -239,6 +243,10 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None: assert "VLLM_CI_HF_TOKEN" in env_names assert "HF_TOKEN" not in env_names assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT + assert step["commands"][1:] == [ + "python3 .buildkite/cuda/scripts/check_vllm_runtime.py", + "pytest -sv tests/example", + ] def _gpu_limit(step: dict) -> int: diff --git a/tests/core/sched/test_generation_batch_coalescing.py b/tests/core/sched/test_generation_batch_coalescing.py index 489e6a819..3cd5431b8 100644 --- a/tests/core/sched/test_generation_batch_coalescing.py +++ b/tests/core/sched/test_generation_batch_coalescing.py @@ -171,6 +171,7 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock s, _ = scheduler(monkeypatch, [event]) s._generation_max_wait_s = 0 s.waiting = [] + s.kv_holding_waiting = [] s.input_coordinator._async_chunk = True s.input_coordinator.update_request_metadata = mocker.Mock() s.input_coordinator.process_pending_chunks = mocker.Mock() @@ -178,7 +179,11 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock s.input_coordinator.update_request_metadata.assert_called_once_with( s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation" ) - s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set()) + s.input_coordinator.process_pending_chunks.assert_called_once() + waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args + assert list(waiting) == [] + assert running is s.running + assert ready == {"0"} and finished == set() @pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))]) diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 546084509..89e4f2847 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/tests/core/sched/test_generation_scheduler_restore.py @@ -11,7 +11,9 @@ from types import SimpleNamespace import pytest import torch from vllm import SamplingParams +from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector from vllm.v1.core.sched.interface import PauseState +from vllm.v1.core.sched.output import CachedRequestData from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue from vllm.v1.request import Request, RequestStatus @@ -53,12 +55,14 @@ class FakeAdapter: def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False): scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) scheduler.max_num_scheduled_tokens = 8 + scheduler.max_num_active_reqs = 1 scheduler.max_num_running_reqs = 1 scheduler._pause_state = PauseState.UNPAUSED scheduler.running = [] scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) scheduler.waiting.add_request(waiting_request) - scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() scheduler.requests = {waiting_request.request_id: waiting_request} scheduler.policy = SchedulingPolicy.FCFS scheduler.chunk_transfer_adapter = FakeAdapter() @@ -86,6 +90,7 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False): get_manager_metadata=lambda: None, ) scheduler.connector = None + scheduler.aux_output_connector = None scheduler.ec_connector = None scheduler.prev_step_scheduled_req_ids = set() scheduler._pending_finish_reqs = [] @@ -105,6 +110,45 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False): return scheduler +def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None: + request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None) + request.block_hashes.append(b"block-hash") + scheduler = _make_generation_scheduler(request, use_v2_model_runner=True) + scheduler.aux_output_connector = AuxOutputSchedulerConnector() + + output = scheduler._build_generation_scheduler_output( + new_reqs_data=[], + cached_reqs_data=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 1}, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0], + ) + + metadata = output.aux_output_connector_metadata + assert metadata is not None + assert metadata.requests == {request.request_id: 0} + assert list(metadata.block_hashes[request.request_id]) == request.block_hashes + + +@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)]) +def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None: + fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None) + holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None) + holding.num_computed_tokens = 1 + scheduler = _make_generation_scheduler(fresh) + scheduler.max_num_running_reqs = runner_slots + scheduler.max_num_active_reqs = active_limit + scheduler.requests[holding.request_id] = holding + scheduler.kv_holding_waiting.add_request(holding) + + output = scheduler.schedule() + + assert set(output.num_scheduled_tokens) == {holding.request_id} + assert list(scheduler.waiting) == [fresh] + assert list(scheduler.kv_holding_waiting) == [] + + def _chunk_request(request_id, **kwargs): defaults = dict( request_id=request_id, @@ -123,7 +167,13 @@ def _chunk_request(request_id, **kwargs): record_event=lambda *args, **kwargs: None, ) defaults.update(kwargs) - return SimpleNamespace(**defaults) + return _HashableChunkRequest(**defaults) + + +class _HashableChunkRequest(SimpleNamespace): + # Deferred queues require identity hashing; SimpleNamespace types it as None. + __hash__ = object.__hash__ # type: ignore[assignment] + __eq__ = object.__eq__ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None: @@ -138,7 +188,8 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> ) scheduler.running = [] scheduler.waiting = create_request_queue(scheduler.policy) - scheduler.skipped_waiting = create_request_queue(scheduler.policy) + scheduler.kv_holding_waiting = create_request_queue(scheduler.policy) + scheduler.deferred_waiting = set() completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None) completed.status = RequestStatus.RUNNING @@ -292,6 +343,7 @@ def _express_scheduler(): scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True) scheduler.waiting.add_request(first) scheduler.requests[first.request_id] = first + scheduler.max_num_active_reqs = 4 scheduler.max_num_running_reqs = 4 scheduler._native_data_plane = True scheduler.chunk_transfer_adapter = None diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py index f8b9d9b4c..3b4e9c702 100644 --- a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py +++ b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py @@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises admits in a finished state -- which kills the stage's engine core, not just the request. -Regression for the sweep missing ``skipped_waiting``, the third queue +Regression for the sweep missing ``kv_holding_waiting``, the third queue upstream admits from. An aborted duplex session parked there was re-selected by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the stage. @@ -49,28 +49,30 @@ class _StubRequest: return (self.priority, self.arrival_time) < (other.priority, other.arrival_time) -def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()): +def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()): """Bare scheduler with just the surface the sweep reads.""" scheduler = OmniARScheduler.__new__(OmniARScheduler) scheduler.waiting = create_request_queue(policy) for req in waiting: scheduler.waiting.add_request(req) - scheduler.skipped_waiting = create_request_queue(policy) - for req in skipped_waiting: - scheduler.skipped_waiting.add_request(req) + scheduler.kv_holding_waiting = create_request_queue(policy) + scheduler.deferred_waiting = set() + for req in kv_holding_waiting: + scheduler.kv_holding_waiting.add_request(req) scheduler.running = list(running) return scheduler @pytest.mark.parametrize("policy", _POLICIES) -def test_schedule_sweeps_skipped_waiting_before_upstream_selection( +def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection( monkeypatch: pytest.MonkeyPatch, policy: SchedulingPolicy, ) -> None: """The public schedule path must remove aborted skipped requests first.""" aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED) - scheduler = _make_scheduler(policy, skipped_waiting=[aborted]) + scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted]) scheduler.chunk_transfer_adapter = None + scheduler.max_num_active_reqs = 8 scheduler.max_num_running_reqs = 8 scheduler.input_coordinator = None scheduler.num_waiting_for_streaming_input = 0 @@ -85,7 +87,7 @@ def test_schedule_sweeps_skipped_waiting_before_upstream_selection( sentinel = object() def upstream_schedule(_self, _throttle_prefills: bool = False): - assert list(scheduler.skipped_waiting) == [] + assert list(scheduler.kv_holding_waiting) == [] return sentinel monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule) diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py index 86102c451..a2ffb1b0a 100644 --- a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py +++ b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py @@ -185,10 +185,15 @@ def test_stage_zero_final_skips_inter_stage_output(mocker): def test_streaming_session_update_invalidates_omits_kv_transfer_cache(): scheduler = OmniARScheduler.__new__(OmniARScheduler) scheduler._omits_kv_transfer_cache = {} - scheduler.skipped_waiting = [] + scheduler.deferred_waiting = set() + scheduler.kv_holding_waiting = [] scheduler.log_stats = False tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0) - session = SimpleNamespace( + + class Session(SimpleNamespace): + __hash__ = object.__hash__ + + session = Session( request_id="req", additional_information=tagged.additional_information, status=RequestStatus.WAITING, diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 747be2ed1..d89cfe7ed 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -11,12 +11,14 @@ from types import MethodType, SimpleNamespace import numpy as np import pytest +import torch from vllm.v1.engine import FinishReason from vllm.v1.outputs import LogprobsLists from vllm.v1.request import RequestStatus from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin +from vllm_omni.outputs import OmniModelRunnerOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -132,11 +134,13 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: scheduler = SimpleNamespace( perf_metrics=None, connector=None, + aux_output_connector=None, chunk_transfer_adapter=None, requests={request.request_id: request for request in requests}, running=list(requests), waiting=_RequestQueue(), - skipped_waiting=_RequestQueue(), + kv_holding_waiting=_RequestQueue(), + deferred_waiting=set(), structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True), transfer_triggered_requests=set(), active_kv_transfers=set(), @@ -292,6 +296,36 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None: np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8]) +@pytest.mark.parametrize("scored", [False, True]) +def test_fixed_token_scores_reach_engine_output(scored) -> None: + request = _Request("req") + request.sampling_params.num_logprobs = None + scheduler = _make_scheduler_stub([request]) + _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False)) + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 3}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]]) + model_runner_output = OmniModelRunnerOutput( + req_ids=["req"], + req_id_to_index={"req": 0}, + sampled_token_ids=[[7]], + prompt_logprobs_dict={}, + prompt_token_id_logprobs_dict={"req": scores} if scored else {}, + ) + + outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output) + (output,) = outputs[0].outputs + + assert output.new_token_ids == [7] + if scored: + torch.testing.assert_close(output.prompt_token_id_logprobs, scores) + else: + assert output.prompt_token_id_logprobs is None + + def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace: return SimpleNamespace( sampled_token_ids=[[]], 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 e5bd8dd4e..042a039fc 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -73,7 +73,8 @@ def _replace_streaming_session(session: Request) -> None: sched.num_waiting_for_streaming_input = 0 sched.log_stats = False sched.chunk_transfer_adapter = None - sched.skipped_waiting = set() + sched.kv_holding_waiting = set() + sched.deferred_waiting = set() session.status = RequestStatus.WAITING_FOR_STREAMING_REQ update = StreamingUpdate( mm_features=None, diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 45de5c05f..785014ff5 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -43,7 +43,8 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS sched.num_waiting_for_streaming_input = 0 sched.log_stats = False sched.chunk_transfer_adapter = None - sched.skipped_waiting = set() + sched.kv_holding_waiting = set() + sched.deferred_waiting = set() sched._free_request_blocks = MagicMock() sched.encoder_cache_manager = MagicMock() sched._inflight_prefills = set() @@ -1185,7 +1186,8 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None: sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex")) sched.running = [] sched.waiting = queue() - sched.skipped_waiting = queue(session) + sched.kv_holding_waiting = queue(session) + sched.deferred_waiting = set() sched.num_waiting_for_streaming_input = 1 sched.chunk_transfer_adapter = SimpleNamespace( receives_chunks=True, @@ -1197,7 +1199,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None: assert session.status == RequestStatus.WAITING assert session in sched.waiting.requests - assert session not in sched.skipped_waiting.requests + assert session not in sched.kv_holding_waiting.requests assert sched.num_waiting_for_streaming_input == 0 assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests @@ -1224,7 +1226,8 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked( receives_chunks=receives_chunks, segment_finished_requests={session.request_id}, ) - sched.skipped_waiting = MagicMock() + sched.kv_holding_waiting = MagicMock() + sched.deferred_waiting = set() sched._enqueue_waiting_request = MagicMock() sched._resume_downstream_chunk_receiver(session) @@ -1232,7 +1235,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked( assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ assert sched.num_waiting_for_streaming_input == 1 assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests - sched.skipped_waiting.remove_requests.assert_not_called() + sched.kv_holding_waiting.remove_requests.assert_not_called() sched._enqueue_waiting_request.assert_not_called() @@ -1288,7 +1291,8 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None: session.num_computed_tokens = 6 session.num_output_placeholders = 0 session.status = RequestStatus.WAITING_FOR_STREAMING_REQ - sched.skipped_waiting.add_request(session) + sched.kv_holding_waiting.add_request(session) + sched.deferred_waiting.add(session) sched.num_waiting_for_streaming_input = 1 @@ -1311,7 +1315,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi assert session.status == RequestStatus.FINISHED_ERROR assert session.request_id not in sched.requests assert len(sched.waiting) == 0 - assert len(sched.skipped_waiting) == 0 + assert len(sched.kv_holding_waiting) == 0 assert sched.num_waiting_for_streaming_input == 0 sched._free_request_blocks.assert_called_once_with(session) assert sched.finished_req_ids == {session.request_id} @@ -1419,7 +1423,8 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler: sched.max_model_len = max_model_len sched.policy = SchedulingPolicy.FCFS sched.waiting = create_request_queue(sched.policy) - sched.skipped_waiting = create_request_queue(sched.policy) + sched.kv_holding_waiting = create_request_queue(sched.policy) + sched.deferred_waiting = set() sched.running = [] sched.requests = {} sched._free_request = MagicMock() # needs a KV manager; not what this covers @@ -1448,7 +1453,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() - assert finished is True assert len(sched.waiting) == 0 - assert len(sched.skipped_waiting) == 0 + assert len(sched.kv_holding_waiting) == 0 assert session.status == RequestStatus.FINISHED_ERROR assert session.resumable is False assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1] @@ -1467,7 +1472,8 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None: finished = sched._handle_stopped_request(session) assert finished is False - assert len(sched.waiting) == 1 + assert list(sched.kv_holding_waiting) == [session] + assert session not in sched.deferred_waiting assert session.status == RequestStatus.WAITING assert not getattr(sched, "_streaming_context_overflow", {}) sched._free_request.assert_not_called() @@ -1489,7 +1495,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None: assert finished is True assert len(sched.waiting) == 0 - assert len(sched.skipped_waiting) == 0 + assert len(sched.kv_holding_waiting) == 0 assert session.status == RequestStatus.FINISHED_ERROR assert session.resumable is False sched._free_request.assert_not_called() @@ -1512,7 +1518,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc finished = sched._handle_stopped_request(session) assert finished is True - assert len(sched.skipped_waiting) == 0 + assert len(sched.kv_holding_waiting) == 0 assert sched.num_waiting_for_streaming_input == 0 assert session.status == RequestStatus.FINISHED_ERROR sched._free_request.assert_not_called() @@ -1545,6 +1551,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati sched.running = [] sched._native_data_plane = native sched.use_v2_model_runner = native + sched.max_num_active_reqs = 8 sched.max_num_running_reqs = 8 sched.input_coordinator = ( SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) @@ -1571,7 +1578,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati observed_limits: list[int] = [] def fake_schedule(self, _throttle_prefills=False): - observed_limits.append(self.max_num_running_reqs) + observed_limits.append(self.max_num_active_reqs) return SimpleNamespace(scheduled_new_reqs=[]) monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule) @@ -1579,4 +1586,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati sched.schedule() assert observed_limits == [7 if native else 8] - assert sched.max_num_running_reqs == 8 + assert sched.max_num_active_reqs == 8 diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index b55a668dd..a2b8f82b7 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -79,16 +79,19 @@ def _make_scheduler( else: hf_config.voxcpm2_runtime_config = runtime_config sched.vllm_config = SimpleNamespace(model_config=model_config) - # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the + # Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the # abort sweep at the top of ``schedule()`` reads it. - sched.skipped_waiting = _MockQueue() + sched.kv_holding_waiting = _MockQueue() + sched.deferred_waiting = set() return sched -def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None: +@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"]) +def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None: scheduler = _make_scheduler() scheduler.running = [_MockRequest("decode")] - scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + scheduler.waiting = _MockQueue() + setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])) assert scheduler._should_defer_waiting_for_unified_decode_graph() @@ -144,12 +147,17 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None: scheduler = _make_scheduler() + scheduler.max_num_active_reqs = 8 scheduler.max_num_running_reqs = 8 scheduler.running = [_MockRequest("decode")] original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING) deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING) original_waiting = _MockQueue([original_waiting_req]) scheduler.waiting = original_waiting + kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING) + original_kv_holding = _MockQueue([kv_holder]) + scheduler.kv_holding_waiting = original_kv_holding + preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED) scheduler.policy = "fcfs" scheduler.chunk_transfer_adapter = None scheduler.input_coordinator = None @@ -160,7 +168,10 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No def fake_upstream_schedule(self, throttle_prefills: bool = False): assert self.waiting is not original_waiting assert not self.waiting + assert self.kv_holding_waiting is not original_kv_holding + assert not self.kv_holding_waiting self.waiting.add_request(deferred_by_upstream) + self.kv_holding_waiting.add_request(preempted_by_upstream) raise RuntimeError("stop before output wrapping") monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule) @@ -169,3 +180,6 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No scheduler.schedule() assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req] + assert scheduler.waiting is original_waiting + assert scheduler.kv_holding_waiting is original_kv_holding + assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder] diff --git a/tests/core/sched/test_omni_cfg_ar_scheduler.py b/tests/core/sched/test_omni_cfg_ar_scheduler.py index e9ff8ea53..b6db7a186 100644 --- a/tests/core/sched/test_omni_cfg_ar_scheduler.py +++ b/tests/core/sched/test_omni_cfg_ar_scheduler.py @@ -23,10 +23,12 @@ def _scheduler(monkeypatch, waiting, running=()): scheduler.waiting = create_request_queue(scheduler.policy) for request in waiting: scheduler.waiting.add_request(request) - scheduler.skipped_waiting = create_request_queue(scheduler.policy) + scheduler.kv_holding_waiting = create_request_queue(scheduler.policy) + scheduler.deferred_waiting = set() scheduler.running = list(running) scheduler.requests = {request.request_id: request for request in [*waiting, *running]} scheduler.max_model_len = 128 + scheduler.max_num_active_reqs = 2 scheduler.max_num_running_reqs = 2 scheduler.max_num_scheduled_tokens = 256 scheduler.block_size = 16 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 0f509c5d2..f1c099622 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -52,10 +52,11 @@ class _SchedulerStub(OmniGenerationScheduler): self.num_waiting_for_streaming_input = 0 self.log_stats = log_stats self.chunk_transfer_adapter = _ChunkTransferAdapterStub() - self.skipped_waiting = _SkippedWaitingStub() + self.kv_holding_waiting = _SkippedWaitingStub() + self.deferred_waiting: set[Request] = set() def _enqueue_waiting_request(self, session: Request) -> None: - raise AssertionError("unexpected enqueue for skipped_waiting miss") + raise AssertionError("unexpected enqueue for kv_holding_waiting miss") class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler): @@ -134,7 +135,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None: sched._handle_stopped_request.side_effect = rearm_resumable_request sched.running = [session] sched.waiting = MagicMock() - sched.skipped_waiting = MagicMock() + sched.kv_holding_waiting = MagicMock() + sched.deferred_waiting = set() sched.structured_output_manager.accept_tokens.return_value = True # Async scheduling can observe the segment boundary in the next schedule() # before this model output is applied and defer the same request for finish. @@ -206,7 +208,8 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete: OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted) ) sched.waiting = MagicMock() - sched.skipped_waiting = MagicMock() + sched.kv_holding_waiting = MagicMock() + sched.deferred_waiting = set() sched.structured_output_manager.accept_tokens.return_value = False sched._pending_finish_reqs = [] sched.recompute_kv_load_failures = False @@ -350,6 +353,8 @@ class _RealignSchedulerStub(OmniSchedulerMixin): self.requests = requests self.running = running self.waiting = waiting + self.kv_holding_waiting: list[Request] = [] + self.deferred_waiting: set[Request] = set() class TestRealignRequestStatusToQueues: diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index c353ced3f..db1f9fb1c 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -123,6 +123,7 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request sched = MagicMock() sched.max_num_scheduled_tokens = 8 + sched.max_num_active_reqs = 8 sched.max_num_running_reqs = 8 sched._native_data_plane = False sched._retains_state_across_chunks = False @@ -135,6 +136,8 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request sched.requests = {request.request_id: request} sched.running = [request] sched.waiting = [] # falsy: the waiting loop never runs + sched.kv_holding_waiting = [] + sched.deferred_waiting = set() sched.policy = SchedulingPolicy.FCFS sched.chunk_transfer_adapter = None sched.scheduler_config = MagicMock(enable_chunked_prefill=True) diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py index 770c85578..9fc0a4c0e 100644 --- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py +++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Integration tests for the post-finish defensive purge of ``self.running`` placed in ``OmniARScheduler.finish_requests`` / @@ -52,6 +55,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting): scheduler.requests = requests scheduler.running = running scheduler.waiting = waiting + scheduler.kv_holding_waiting = [] + scheduler.deferred_waiting = set() return scheduler diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index 2ecd2772a..e1f9bc8c9 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -52,6 +52,8 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids( scheduler.requests = {} scheduler.running = [] scheduler.waiting = [] + scheduler.kv_holding_waiting = [] + scheduler.deferred_waiting = set() def fake_finish_requests(self, request_ids, finished_status): assert request_ids == ["req-a", "req-b"] diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index 82956e27a..086593593 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py @@ -7,7 +7,10 @@ from typing import Any import pytest from vllm.config import SchedulerConfig, VllmConfig +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.engine import FinishReason +from vllm.v1.request import Request from vllm_omni.config.model import OmniModelConfig from vllm_omni.core.sched import omni_scheduler_mixin @@ -18,7 +21,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _Scheduler(OmniSchedulerMixin): - pass + _is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status) def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch): @@ -120,7 +123,11 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): calls: list[tuple[Any, ...]] = [] scheduler = _Scheduler() - scheduler.waiting = ["waiting"] + waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None) + holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None) + holding.num_computed_tokens = 1 + scheduler.waiting = [waiting] + scheduler.kv_holding_waiting = [holding] scheduler.running = ["running"] scheduler.requests = {"request": object()} scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode)) @@ -137,16 +144,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): scheduler.chunk_transfer_adapter = SimpleNamespace( receives_chunks=True, process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append( - ("process", waiting, running, scheduler_requests) + ("process", list(waiting), running, scheduler_requests) ), restore_queues=lambda waiting, running, scheduler_requests: calls.append( - ("restore-chunks", waiting, running, scheduler_requests) + ("restore-chunks", list(waiting), running, scheduler_requests) ), collect_timed_out_request_ids=_collect_timed_out, collect_failed_send_request_ids=_collect_failed_sends, ) scheduler.input_coordinator = SimpleNamespace( - restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running)) + restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running)) ) scheduler._process_pending_omni_inputs("ar") @@ -155,13 +162,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): assert calls == [ ("consume", "ar"), ("timeouts",), - ("process", scheduler.waiting, scheduler.running, scheduler.requests), + ("process", [holding, waiting], scheduler.running, scheduler.requests), # The chunk deadline runs after chunks are applied, so a chunk that # arrived this cycle resets the clock before it is measured (R1.1). ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S), ("failed-sends",), - ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests), - ("restore-full", scheduler.waiting, scheduler.running), + ("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests), + ("restore-full", [holding, waiting], scheduler.running), ] diff --git a/tests/core/sched/test_omni_scheduler_ready_inbox.py b/tests/core/sched/test_omni_scheduler_ready_inbox.py index 483255f21..db215eeb0 100644 --- a/tests/core/sched/test_omni_scheduler_ready_inbox.py +++ b/tests/core/sched/test_omni_scheduler_ready_inbox.py @@ -5,8 +5,13 @@ from types import SimpleNamespace import pytest +from vllm import SamplingParams +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import Request, RequestStatus +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin +from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator from vllm_omni.outputs import OmniConnectorOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -15,6 +20,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker): scheduler = OmniSchedulerMixin() scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], [] + scheduler.kv_holding_waiting = [] + scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status coordinator = SimpleNamespace( _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock() ) @@ -32,7 +39,92 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker): coordinator.update_request_metadata.assert_called_once_with( scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar" ) - coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"}) + coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"}) + assert list(coordinator.process_pending_chunks.call_args.args[0]) == [] scheduler._consume_pending_connector_output(model_mode="ar") assert coordinator.update_request_metadata.call_count == 1 - coordinator.process_pending_chunks.assert_called_with([], [], set(), set()) + coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set()) + assert list(coordinator.process_pending_chunks.call_args.args[0]) == [] + + +@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY]) +@pytest.mark.parametrize("async_chunk", [False, True]) +def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(policy) + scheduler.kv_holding_waiting = create_request_queue(policy) + scheduler.deferred_waiting = set() + scheduler.running = [] + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = OmniSchedulingCoordinator( + stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk + ) + scheduler._init_omni_connector_output_inbox() + holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None) + holder.num_computed_tokens = 2 + fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None) + scheduler.requests = {request.request_id: request for request in (holder, fresh)} + scheduler._enqueue_waiting_request(fresh) + scheduler._enqueue_waiting_request(holder) + + scheduler._consume_pending_connector_output(model_mode="ar") + + assert not scheduler.waiting and not scheduler.kv_holding_waiting + expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT + assert holder.status == fresh.status == expected_status + coordinator = scheduler.input_coordinator + registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations + assert [entry.request_id for entry in registrations] == ["holder", "fresh"] + + scheduler._restore_omni_wait_queues() + assert list(scheduler.kv_holding_waiting) == [holder] + assert list(scheduler.waiting) == [fresh] + ready_output = ( + OmniConnectorOutput(chunk_ready_req_ids={"holder"}) + if async_chunk + else OmniConnectorOutput(stage_recv_req_ids={"holder"}) + ) + scheduler.enqueue_omni_connector_output(ready_output) + scheduler._consume_pending_connector_output(model_mode="ar") + + assert holder.status == RequestStatus.WAITING + assert fresh.status == expected_status + assert list(scheduler.kv_holding_waiting) == [holder] + assert not scheduler.waiting + scheduler._restore_omni_wait_queues() + assert list(scheduler.kv_holding_waiting) == [holder] + assert list(scheduler.waiting) == [fresh] + + +@pytest.mark.parametrize("async_chunk", [False, True]) +@pytest.mark.parametrize( + "status", + [ + RequestStatus.WAITING_FOR_STREAMING_REQ, + RequestStatus.WAITING_FOR_REMOTE_KVS, + RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR, + ], +) +def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.running = [] + request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None) + request.status = status + request.num_computed_tokens = 2 + scheduler.requests = {request.request_id: request} + scheduler._enqueue_waiting_request(request) + scheduler.input_coordinator = OmniSchedulingCoordinator( + stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk + ) + scheduler._init_omni_connector_output_inbox() + + scheduler._consume_pending_connector_output(model_mode="ar") + + assert request.status == status + assert list(scheduler.kv_holding_waiting) == [request] + assert scheduler.deferred_waiting == {request} + assert not scheduler.input_coordinator.pending_chunk_registrations + assert not scheduler.input_coordinator.pending_input_registrations diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py index bc392abca..d58b9f9a4 100644 --- a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py +++ b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py @@ -59,7 +59,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp scheduler.requests = requests scheduler.running = running scheduler.waiting = waiting - scheduler.skipped_waiting = skipped if skipped is not None else [] + scheduler.kv_holding_waiting = skipped if skipped is not None else [] + scheduler.deferred_waiting = set() scheduler.num_waiting_for_streaming_input = counter return scheduler @@ -97,8 +98,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status( assert scheduler.num_waiting_for_streaming_input == 0 -def test_resync_counts_parked_requests_in_skipped_waiting() -> None: - """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``, +def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None: + """``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``, so the resync has to look at both or it under-counts and reintroduces the spin the counter exists to prevent.""" parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ) @@ -128,9 +129,11 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk scheduler.requests = {parked.request_id: parked} scheduler.running = [parked] scheduler.waiting = [] - scheduler.skipped_waiting = [] + scheduler.kv_holding_waiting = [] + scheduler.deferred_waiting = set() scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite scheduler.policy = "fcfs" + scheduler.max_num_active_reqs = 8 scheduler.max_num_running_reqs = 8 seen: dict[str, int] = {} @@ -161,12 +164,13 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() - scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) scheduler.waiting.add_request(aborted) - scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS) - scheduler.skipped_waiting.add_request(aborted) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.kv_holding_waiting.add_request(aborted) scheduler.running = [aborted] scheduler._drop_aborted_queued_requests() assert list(scheduler.waiting) == [] - assert list(scheduler.skipped_waiting) == [] + assert list(scheduler.kv_holding_waiting) == [] assert scheduler.running == [] diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py index a6d5fd342..566f8ab43 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py @@ -216,44 +216,40 @@ def test_video_chat_stream_invalid_requests( # WS /v1/realtime # ───────────────────────────────────────────────────────────────────────────── -# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake). -_REALTIME_WS_MODEL_MISMATCH = object() -_REALTIME_WS_INVALID_AUDIO_APPEND = object() - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize( "send_frames_spec, ws_error_code, err_message", [ - pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"), + pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"), pytest.param( json.dumps({"type": "session.update"}), "invalid_event", - ("model", "Missing required field"), - id="session_update_missing_model", + "Invalid or unrecognized client event", + id="session_update_missing_session", ), pytest.param( - _REALTIME_WS_MODEL_MISMATCH, - "model_not_found", - ("model", "does not exist"), - id="session_update_model_mismatch", + json.dumps({"type": "session.update", "session": []}), + "invalid_event", + "Invalid or unrecognized client event", + id="session_update_invalid_session", ), pytest.param( - json.dumps({"type": "input_audio_buffer.commit", "final": False}), - "model_not_validated", - ("session.update", "validate the model"), - id="commit_without_model_validation", + json.dumps({"type": "input_audio_buffer.commit"}), + "invalid_request_error", + "Input audio buffer is empty", + id="commit_without_audio", ), pytest.param( json.dumps({"type": "unknown.event"}), - "unknown_event", - "Unknown event type", + "invalid_event", + "Invalid or unrecognized client event", id="unknown_event_type", ), pytest.param( - _REALTIME_WS_INVALID_AUDIO_APPEND, - "invalid_audio", - "Invalid audio data", + json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}), + "invalid_request_error", + "Invalid base64 audio data", id="invalid_audio_append", ), ], @@ -266,27 +262,12 @@ def test_realtime_invalid_requests( ws_error_code: str, err_message: str | tuple[str, ...], ) -> None: - send_frames: Any - if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH: - send_frames = json.dumps( - { - "type": "session.update", - "model": "this-model-is-not-served-on-this-instance", - } - ) - elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND: - send_frames = [ - json.dumps({"type": "session.update", "model": omni_server.model}), - json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}), - ] - else: - send_frames = send_frames_spec online_client.send_realtime_ws_request( { - "send_frames": send_frames, + "send_frames": send_frames_spec, "timeout": 120, "ws_max_size": None, - "ws_skip_types": ["session.created"], + "ws_skip_types": ["session.created", "conversation.created"], "ws_json_type": "error", "ws_error_code": ws_error_code, "err_message": err_message, diff --git a/tests/diffusion/cache/test_seacache.py b/tests/diffusion/cache/test_seacache.py index 5f88db5c6..a876b6b36 100644 --- a/tests/diffusion/cache/test_seacache.py +++ b/tests/diffusion/cache/test_seacache.py @@ -329,7 +329,9 @@ def _uneven_cfg_sharded_worker( init_method=init_method, rank=rank, world_size=4, - timeout=datetime.timedelta(seconds=10), + # Spawned workers import torch/vLLM before joining; allow startup + # skew without changing the collective deadlock checks below. + timeout=datetime.timedelta(seconds=60), ) try: fs_groups = [ diff --git a/tests/diffusion/diffusion_kv/test_kv_cache_utils.py b/tests/diffusion/diffusion_kv/test_kv_cache_utils.py index 543524404..ba617079f 100644 --- a/tests/diffusion/diffusion_kv/test_kv_cache_utils.py +++ b/tests/diffusion/diffusion_kv/test_kv_cache_utils.py @@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch): # Independent expected native extras: a crossing item has a negative # relative offset; a block can include the end of one item and the next. first, second = (feature.identifier for feature in row.mm_features) - extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)] + extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)] parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace))) expected = [] for start, keys in zip(range(0, 12, 4), extras, strict=True): diff --git a/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py b/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py index 2d3ae0a20..f1974d326 100644 --- a/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py +++ b/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py @@ -31,6 +31,10 @@ def producer(monkeypatch): # vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34, # #56033), the MLA/Mamba flags and the head-resharding validator. worker.pp_size = 1 + worker.opaque_packed_storages = set() + worker._prepared_transfer_regions = {} + worker.region_shared_groups = [] + worker.region_row_offsets = [] worker.use_mla = False worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False) worker._validate_head_resharding_layout = Mock(return_value=None) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 49aba9140..8eb60202c 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -1253,14 +1253,15 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads speculative_config=None, ) - def select_backend(*, backend, attn_selector_config, num_heads): - # Exercise the real 0.29 selector, replacing only platform resolution. + def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype): + # Exercise the real upstream selector, replacing only platform resolution. # Ulysses has already gathered tokens and sharded heads before the # native kernel; the MoE PCP mapping must not request PCP attention. assert not attn_selector_config.use_pcp assert not attn_selector_config.use_dcp assert attn_selector_config.use_non_causal assert num_heads == 4 + assert _run_kv_cache_dtype == config.cache_config.cache_dtype selected_backends.append((backend, config.attention_config.backend_per_kind)) if selection_fails: raise ValueError("test backend unavailable") diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 21cb18821..3c4f9ba03 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -14,6 +14,7 @@ from typing import Literal import pytest import torch +from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D from vllm_omni.diffusion.distributed.parallel_state import ( @@ -56,6 +57,8 @@ def _init_worker( "MASTER_PORT": str(master_port), } ) + if device_kind == "cpu": + configure_cpu_collective_worker() backend = "gloo" if device_kind == "cpu" else None init_distributed_environment(backend=backend) @@ -227,6 +230,10 @@ def _run_ring_p2p( atol=atol, msg=f"[Rank {local_rank}] Ring P2P data mismatch", ) + if device_kind == "cpu": + # P2P completion only synchronizes neighbors. Keep fast CPU ranks + # alive until every rank has finished constructing and using groups. + torch.distributed.barrier() finally: destroy_distributed_env() @@ -243,7 +250,7 @@ def _spawn_4d_identity( device_kind: DeviceKind, master_port: int, ) -> None: - torch.multiprocessing.spawn( + start_collective_workers( _run_4d_identity, args=( world_size, @@ -257,6 +264,7 @@ def _spawn_4d_identity( master_port, ), nprocs=world_size, + device_kind=device_kind, ) @@ -272,7 +280,7 @@ def _spawn_5d_identity( device_kind: DeviceKind, master_port: int, ) -> None: - torch.multiprocessing.spawn( + start_collective_workers( _run_5d_identity, args=( world_size, @@ -286,6 +294,7 @@ def _spawn_5d_identity( master_port, ), nprocs=world_size, + device_kind=device_kind, ) @@ -299,10 +308,11 @@ def _spawn_ring_p2p( device_kind: DeviceKind, master_port: int, ) -> None: - torch.multiprocessing.spawn( + start_collective_workers( _run_ring_p2p, args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port), nprocs=world_size, + device_kind=device_kind, ) @@ -499,8 +509,9 @@ def _run_fused_qkv_matches_three_4d( @pytest.mark.cpu @pytest.mark.parametrize("world_size", [2, 4]) def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None: - torch.multiprocessing.spawn( + start_collective_workers( _run_fused_qkv_matches_three_4d, args=(world_size, "cpu", 29613), nprocs=world_size, + device_kind="cpu", ) diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 7cebbcc53..c64e880cb 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -19,6 +19,7 @@ from vllm.utils.network_utils import get_file_store_init_method from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module +from tests.helpers.distributed import collective_context, start_collective_workers from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( @@ -39,6 +40,25 @@ DeviceKind = Literal["cpu", "cuda"] _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +@pytest.fixture(autouse=True) +def _restore_collective_platforms(): + """CPU baselines run in the parent process; retain its original platforms.""" + import vllm.platforms + + from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state + + original_platform = vllm.platforms.current_platform + original_omni_platforms = { + module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state) + } + try: + yield + finally: + vllm.platforms.current_platform = original_platform + for module, platform in original_omni_platforms.items(): + module.current_omni_platform = platform + + def _cleanup_distributed() -> None: """Destroy omni distributed state left dangling after a test. @@ -409,6 +429,10 @@ def init_dist( device = _worker_device(local_rank, device_kind) if device_kind == "cuda": current_omni_platform.set_device(device) + else: + from tests.helpers.distributed import configure_cpu_collective_worker + + configure_cpu_collective_worker() backend = "gloo" if device_kind == "cpu" else None init_distributed_environment( world_size=world_size, @@ -500,13 +524,14 @@ def isend_irecv_worker( def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None: - mp_context = torch.multiprocessing.get_context("spawn") + mp_context = collective_context(device_kind) manager = mp_context.Manager() q = manager.Queue() - torch.multiprocessing.spawn( + start_collective_workers( isend_irecv_worker, args=(pp_size, init_method, device_kind, q), nprocs=pp_size, + device_kind=device_kind, ) results = {label: tensor for label, tensor in [q.get(), q.get()]} torch.testing.assert_close( @@ -656,15 +681,16 @@ def _run_predict_noise( baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind) - mp_context = torch.multiprocessing.get_context("spawn") + mp_context = collective_context(device_kind) manager = mp_context.Manager() pp_q = manager.Queue() world_size = pp_size * cfg_size - torch.multiprocessing.spawn( + start_collective_workers( predict_noise_worker, args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), nprocs=world_size, + device_kind=device_kind, ) pp_out = pp_q.get() @@ -900,15 +926,16 @@ def _run_scheduler_step( baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind) - mp_context = torch.multiprocessing.get_context("spawn") + mp_context = collective_context(device_kind) manager = mp_context.Manager() q = manager.Queue() world_size = pp_size * cfg_size - torch.multiprocessing.spawn( + start_collective_workers( scheduler_step_worker, args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), nprocs=world_size, + device_kind=device_kind, ) result = q.get() diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index b2dea0c44..4065cc517 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo ) def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm): from transformers import AutoModel - from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + from vllm_omni.quantization.fp8_config import DiffusionFp8Config od_config = OmniDiffusionConfig( model="dummy-boogu", @@ -287,15 +287,14 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co mock_dependencies["mllm_config_loader"].assert_not_called() mock_dependencies["mllm_builder"].assert_not_called() assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"] - assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config) + assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config) def test_constructor_preserves_mllm_ignored_layers(mock_dependencies): - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + from vllm_omni.quantization.fp8_config import DiffusionFp8Config - quant_config = Fp8Config( + quant_config = DiffusionFp8Config( ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"] ) od_config = OmniDiffusionConfig( diff --git a/tests/diffusion/models/flux2/test_flux2_quantization.py b/tests/diffusion/models/flux2/test_flux2_quantization.py index 3c461acf5..8ba413b9e 100644 --- a/tests/diffusion/models/flux2/test_flux2_quantization.py +++ b/tests/diffusion/models/flux2/test_flux2_quantization.py @@ -16,17 +16,13 @@ class _RecordingBlock(nn.Module): def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch): - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm_omni.diffusion.models.flux2 import flux2_transformer + from vllm_omni.quantization import build_quant_config monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock) monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock) - quant_config = Fp8Config( - is_checkpoint_fp8_serialized=False, - activation_scheme="dynamic", - ) + quant_config = build_quant_config("fp8") model = flux2_transformer.Flux2Transformer2DModel( in_channels=4, out_channels=4, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index e1f83fd01..0e6197c96 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None): def test_fp8_scope_and_prefix_propagation(monkeypatch): from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod - from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + from vllm_omni.quantization.fp8_config import DiffusionFp8Config monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear) monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear) @@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch): "blocks.0.adaln_proj.linear", "final_layer.adaln_proj.linear", } - fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers)) + fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers)) model = h3.MiniMaxH3DiTModel( _small_od_config(), quant_config=fp8_config, diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index 9bad16792..ce0daeb58 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -15,6 +15,74 @@ def test_build_quant_config_fp8(): assert config is not None assert config.get_name() == "fp8" assert config.activation_scheme == "dynamic" + assert not config.is_checkpoint_fp8_serialized + + +def test_build_quant_config_fp8_uses_native_online_quantization(): + from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym + + from vllm_omni.quantization import build_quant_config + + config = build_quant_config("fp8", ignored_layers=["proj_out"]) + assert isinstance(config, OnlineQuantizationConfig) + assert config.args.linear.weight == kFp8StaticTensorSym + assert config.args.moe.weight == kFp8StaticTensorSym + assert config.ignored_layers == ["proj_out"] + + +def test_build_quant_config_upstream_online_fp8_shorthand(): + from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym + + from vllm_omni.quantization import build_quant_config + + config = build_quant_config("fp8_per_tensor", ignore=["proj_out"]) + assert isinstance(config, OnlineQuantizationConfig) + assert config.args.linear.weight == kFp8StaticTensorSym + assert config.ignored_layers == ["proj_out"] + + +def test_online_fp8_rejects_static_activation(): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="activation_scheme='dynamic'"): + build_quant_config("fp8", activation_scheme="static") + + +def test_build_online_quant_config_preserves_explicit_native_args(): + from vllm.config.quantization import QuantizationConfigArgs, QuantSpec + from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + + from vllm_omni.quantization import build_quant_config + + args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"]) + config = build_quant_config("online", args=args) + assert isinstance(config, OnlineQuantizationConfig) + assert config.args is args + assert config.ignored_layers == ["proj_out"] + + +def test_build_online_quant_config_requires_args(): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="requires quantization config arguments"): + build_quant_config("online") + + +def test_serialized_checkpoint_replaces_online_fp8_config(): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + config = resolve_quant_config_from_disk( + build_quant_config("fp8"), + {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}, + ) + assert isinstance(config, Fp8Config) + assert config.is_checkpoint_fp8_serialized + assert config.activation_scheme == "static" def test_build_quant_config_none(): @@ -39,7 +107,7 @@ def test_build_quant_config_invalid(): def test_build_quant_config_dict(): from vllm_omni.quantization import build_quant_config - config = build_quant_config({"method": "fp8", "activation_scheme": "static"}) + config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}) assert config is not None assert config.get_name() == "fp8" assert config.activation_scheme == "static" @@ -48,7 +116,7 @@ def test_build_quant_config_dict(): def test_build_quant_config_dict_not_mutated(): from vllm_omni.quantization import build_quant_config - original = {"method": "fp8", "activation_scheme": "static"} + original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"} copy = original.copy() build_quant_config(original) assert original == copy @@ -99,7 +167,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated(): """Inner component dicts should not be mutated by build_quant_config.""" from vllm_omni.quantization import build_quant_config - inner = {"method": "fp8", "activation_scheme": "static"} + inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"} original = inner.copy() build_quant_config({"transformer": inner, "vae": None}) assert inner == original @@ -115,20 +183,18 @@ def test_flat_dict_not_misdetected_as_per_component(): def test_build_quant_config_passthrough(): - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.fp8_config import DiffusionFp8Config - fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic") + fp8 = DiffusionFp8Config() assert build_quant_config(fp8) is fp8 def test_component_config_routing(): - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm_omni.quantization import ComponentQuantizationConfig + from vllm_omni.quantization.fp8_config import DiffusionFp8Config - fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic") + fp8 = DiffusionFp8Config() config = ComponentQuantizationConfig( component_configs={"transformer": fp8, "vae": None}, ) @@ -140,11 +206,10 @@ def test_component_config_routing(): def test_component_config_with_default(): - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm_omni.quantization import ComponentQuantizationConfig + from vllm_omni.quantization.fp8_config import DiffusionFp8Config - fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic") + fp8 = DiffusionFp8Config() config = ComponentQuantizationConfig( component_configs={"vae": None}, default_config=fp8, @@ -167,7 +232,7 @@ def test_integration_dict(): config = OmniDiffusionConfig( model="test", - quantization_config={"method": "fp8", "activation_scheme": "static"}, + quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}, ) assert config.quantization_config is not None assert config.quantization_config.get_name() == "fp8" diff --git a/tests/diffusion/quantization/test_online_fp8_gpu.py b/tests/diffusion/quantization/test_online_fp8_gpu.py new file mode 100644 index 000000000..e3cccb355 --- /dev/null +++ b/tests/diffusion/quantization/test_online_fp8_gpu.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise online FP8 loading and inference through the native CUDA path.""" + +import pytest +import torch +from torch.nn import functional as F +from vllm.config import set_current_vllm_config +from vllm.distributed import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod +from vllm.utils.network_utils import get_open_port +from vllm.utils.torch_utils import set_default_torch_dtype + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config +from vllm_omni.quantization import build_quant_config + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"]) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method): + ignored_key = "ignored_layers" if method == "fp8" else "ignore" + config = build_quant_config(method, **{ignored_key: ["skip"]}) + od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config) + vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config) + + try: + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + backend="nccl", + ) + with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"): + initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + layer = ReplicatedLinear( + 128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True + ) + assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod) + assert layer.weight.is_meta + + generator = torch.Generator(device="cuda").manual_seed(42) + weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5 + inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16) + expected = F.linear(inputs, weights).float() + + # This is the loader installed by the native online quantizer. + # It materializes and quantizes the BF16 checkpoint in one pass. + layer.weight.weight_loader(layer.weight, weights) + assert layer.weight.dtype == torch.float8_e4m3fn + actual = layer(inputs)[0].float() + relative_mse = (actual - expected).square().mean() / expected.square().mean() + assert relative_mse.item() < 0.002 + + ignored = ReplicatedLinear( + 128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True + ) + assert isinstance(ignored.quant_method, UnquantizedLinearMethod) + assert ignored.weight.dtype == torch.bfloat16 + finally: + destroy_model_parallel() + destroy_distributed_environment() diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 438414854..f0b8d45b5 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -11,6 +11,7 @@ import pytest import torch from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats from vllm.v1.request import Request, RequestStatus @@ -635,6 +636,45 @@ def build_adapter(monkeypatch, mocker: MockerFixture): return _build +@pytest.mark.parametrize("active_stream_window", [0, 2]) +def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window): + adapter, _ = build_adapter(active_stream_window=active_stream_window) + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.running = [] + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None) + holder.num_computed_tokens = 2 + fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None) + scheduler.requests = {request.request_id: request for request in (holder, fresh)} + scheduler._enqueue_waiting_request(fresh) + scheduler._enqueue_waiting_request(holder) + + scheduler._process_pending_omni_inputs(model_mode="ar") + + assert not scheduler.waiting and not scheduler.kv_holding_waiting + assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh] + assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK + scheduler._restore_omni_wait_queues() + assert list(scheduler.kv_holding_waiting) == [holder] + assert list(scheduler.waiting) == [fresh] + + adapter._finished_load_reqs.add(holder.request_id) + scheduler._process_pending_omni_inputs(model_mode="ar") + + assert holder.status == RequestStatus.WAITING + assert fresh.status == RequestStatus.WAITING_FOR_CHUNK + assert list(scheduler.kv_holding_waiting) == [holder] + assert not scheduler.waiting + assert adapter.requests_with_ready_chunks == {holder.request_id} + scheduler._restore_omni_wait_queues() + assert list(scheduler.kv_holding_waiting) == [holder] + assert list(scheduler.waiting) == [fresh] + + def _dequeue_load_entry(adapter, request): """Register and take the entry as the background recv loop would.""" adapter.load_async(request) @@ -2445,7 +2485,8 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity( scheduler.requests = {request.request_id: request} scheduler.running = [request] if placement == "running" else [] scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else []) - scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else []) + scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else []) + scheduler.deferred_waiting = set() scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned) freed = [] @@ -2469,7 +2510,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity( assert not adapter.waiting_for_chunk_running_requests assert scheduler.running == [] assert list(scheduler.waiting) == [] - assert list(scheduler.skipped_waiting) == [] + assert list(scheduler.kv_holding_waiting) == [] assert scheduler.num_waiting_for_streaming_input == 0 fresh_a = _req("req-fresh-a", RequestStatus.WAITING) @@ -2492,7 +2533,8 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt scheduler.requests = {request.request_id: request} scheduler.running = [] scheduler.waiting = DummyWaitingQueue() - scheduler.skipped_waiting = DummyWaitingQueue() + scheduler.kv_holding_waiting = DummyWaitingQueue() + scheduler.deferred_waiting = set() scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed") assert ( @@ -2566,6 +2608,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc scheduler = mocker.MagicMock() scheduler._first_chunk_express = False + scheduler.aux_output_connector = None scheduler._express_min_slack_s = 0.0 scheduler.chunk_transfer_adapter = adapter_mock scheduler.connector = None @@ -2646,6 +2689,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw)) scheduler = mocker.MagicMock() + scheduler.aux_output_connector = None scheduler.chunk_transfer_adapter = adapter_mock bind_omits_transfer_helpers(scheduler) scheduler.connector = None @@ -2744,6 +2788,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): sched.requests = {} sched.running = [] sched.waiting = [] + sched.kv_holding_waiting = [] + sched.deferred_waiting = set() with patch.object(VLLMScheduler, "finish_requests", _super_finish): OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) @@ -2759,7 +2805,8 @@ def _parked_sender_scheduler(adapter, session): scheduler.requests = {session.request_id: session} scheduler.running = [] scheduler.waiting = DummyWaitingQueue() - scheduler.skipped_waiting = DummyWaitingQueue([session]) + scheduler.kv_holding_waiting = DummyWaitingQueue([session]) + scheduler.deferred_waiting = set() scheduler.num_waiting_for_streaming_input = 1 scheduler.finished_req_ids_dict = {} diff --git a/tests/e2e/accuracy/sana_video/conftest.py b/tests/e2e/accuracy/sana_video/conftest.py new file mode 100644 index 000000000..427dae468 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/conftest.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + + +@pytest.fixture +def sana_model_parallel(tmp_path): + """Construct native parallel layers with real single-rank GPU groups.""" + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, + ) + + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=f"file://{tmp_path / 'sana-distributed'}", + ) + try: + initialize_model_parallel() + yield + finally: + destroy_distributed_env() diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py index cb6ffe6c9..87c91d093 100644 --- a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py +++ b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py @@ -202,6 +202,7 @@ def test_sana_video_transformer_matches_frozen_golden( variant: str, _hardware, accuracy_artifact_root: Path, + sana_model_parallel, ) -> None: del _hardware from vllm.config import LoadConfig diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py index aabe166d2..787114529 100644 --- a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py +++ b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py @@ -189,6 +189,7 @@ def _cfg_prediction( def test_sana_video_quick_pipeline_stages_match_diffusers( variant: str, _hardware, + sana_model_parallel, ) -> None: """Compare the first denoising step without replacing native attention. diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index 79460c631..f27bb8da9 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -23,7 +23,8 @@ BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) def _get_diffusion_feature_cases(model: str): return [ pytest.param( - OmniServerParams(model=model), + # The unsharded transformer alone exceeds the L4's 22 GiB. + OmniServerParams(model=model, server_args=["--enable-layerwise-offload"]), id="base", marks=BASE_FEATURE_MARKS, ), diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py index bd38a7fc7..b2c40673a 100644 --- a/tests/e2e/online_serving/test_moss_tts_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -28,7 +28,7 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config pytestmark = [ pytest.mark.advanced_model, @@ -84,7 +84,12 @@ tts_server_params = [ pytest.param( OmniServerParams( model=MODEL, - stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"), + stage_config_path=modify_stage_config( + get_deploy_config_path("moss_tts_realtime.yaml"), + # Leave room on the single L4 for the codec weights, reference + # encoder, and speech-serving activations in addition to KV. + updates={"stages": {0: {"gpu_memory_utilization": 0.35}}}, + ), server_args=["--disable-log-stats"], ), id="moss_tts_realtime", diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py index 389b0b6ba..b5ecd3435 100644 --- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py @@ -21,7 +21,14 @@ models = ["Qwen/Qwen2.5-Omni-7B"] # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the # platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml. -stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))] +stage_configs = [ + modify_stage_config( + get_deploy_config_path("ci/qwen2_5_omni.yaml"), + # This suite requests two cards. Share the Talker's card with + # Code2Wav and leave room for its whole-sequence DiT workspace. + updates={"stages": {1: {"gpu_memory_utilization": 0.3}, 2: {"devices": "1"}}}, + ) +] # Create parameter combinations for model and stage config test_params = [ diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 142affab7..bfa7be0e4 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -20,7 +20,11 @@ import pytest from omegaconf import OmegaConf from pydantic import ValidationError from transformers import PretrainedConfig, Qwen3OmniMoeConfig +from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs +from vllm.exceptions import VLLMValidationError +from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState from tests.helpers.mock import patch_hf_snapshot_download from vllm_omni.config.model import OmniModelConfig @@ -32,6 +36,36 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("use_v2", [False, True]) +def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2): + from vllm.sampling_params import SamplingParams + + from vllm_omni.model_executor.models.audex.cfg import AudexCFGLogitsProcessor + + # Upstream's process-wide preference must not override a stage's deploy choice. + monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", not use_v2) + config = VllmConfig() + model_config = object.__new__(OmniModelConfig) + model_config.use_v2_model_runner = use_v2 + model_config.runner_type = "generate" + model_config.logits_processors = [LogitBiasState if use_v2 else AudexCFGLogitsProcessor] + config.model_config = model_config + assert config.use_v2_model_runner is use_v2 + processor = object.__new__(InputProcessor) + processor.vllm_config, processor.model_config = config, model_config + validator = processor._build_logits_processors_params_validator() + validator(SamplingParams(extra_args={"cfg_role": "cond", "cfg_scale": 1.5})) + if not use_v2: + with pytest.raises(VLLMValidationError, match="cfg_role"): + validator(SamplingParams(extra_args={"cfg_role": "invalid"})) + + +@pytest.mark.parametrize("use_v2", [False, True]) +def test_plain_vllm_config_retains_upstream_runner_preference(monkeypatch, use_v2): + monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", use_v2) + assert VllmConfig().use_v2_model_runner is use_v2 + + def test_sync_config_is_omni(): """Ensure create_model_config gives the right type.""" cfg = OmniEngineArgs().create_model_config() diff --git a/tests/engine/test_arg_utils_shared_fields.py b/tests/engine/test_arg_utils_shared_fields.py index 40a584fce..cdf5f272d 100644 --- a/tests/engine/test_arg_utils_shared_fields.py +++ b/tests/engine/test_arg_utils_shared_fields.py @@ -8,6 +8,7 @@ from __future__ import annotations from dataclasses import fields import pytest +from vllm.config import AuxOutputConfig from vllm_omni.engine.arg_utils import ( SHARED_FIELDS, @@ -49,3 +50,26 @@ def test_internal_blacklist_keys_derived_from_orchestrator(): """ blacklist = internal_blacklist_keys() assert blacklist == orchestrator_field_names() - SHARED_FIELDS + + +@pytest.mark.parametrize( + "aux_args", + [ + {"enable_return_routed_experts": True}, + {"aux_output_config": {"enable_return_routed_experts": True}}, + {"aux_output_config": AuxOutputConfig(enable_return_routed_experts=True)}, + ], +) +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) +@pytest.mark.parametrize("is_npu", [False, True]) +def test_aux_outputs_require_a_supported_omni_stage_runner(aux_args, use_v2_model_runner, is_npu, monkeypatch): + # Upstream runner selection reads this environment variable; Omni workers + # select their runner from the separate per-stage flag instead. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + monkeypatch.setattr("vllm_omni.engine.arg_utils.current_omni_platform.is_npu", lambda: is_npu) + if use_v2_model_runner or is_npu: + args = OmniEngineArgs(model="unused", use_v2_model_runner=use_v2_model_runner, **aux_args) + assert args.aux_output_config.enabled + else: + with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"): + OmniEngineArgs(model="unused", **aux_args) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 7ad111e25..8947d109f 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -168,6 +168,37 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock output_processor.add_request.assert_not_called() +def test_build_add_request_message_preserves_kv_hints_after_omni_payload_upgrade(mocker): + from vllm.v1.engine.input_processor import InputProcessor + from vllm.v1.kv_hints import KvHintsEnvelope + + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.default_sampling_params_list = [params] + engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")] + engine.supported_tasks = ("generate",) + processor = mocker.create_autospec(InputProcessor, instance=True) + hints = KvHintsEnvelope("1", "hint-message", []) + + def process(**kwargs): + request = _make_engine_core_request() + request.kv_hints = kwargs["kv_hints"] + return request + + processor.process_inputs.side_effect = process + engine.input_processor = processor + msg = engine._build_add_request_message( + request_id="req-1", + prompt={"prompt_token_ids": [1], "additional_information": {"speaker": "voice"}}, + sampling_params_list=[params], + kv_hints=hints, + ) + assert processor.process_inputs.call_args.kwargs["kv_hints"] is hints + assert isinstance(msg.prompt, OmniEngineCoreRequest) + assert msg.prompt.kv_hints is hints + assert msg.prompt.additional_information is not None + + def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index a5314ea26..1a92c7b6e 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -296,11 +296,21 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): """Regression test for forward references in make_async_mp_client.""" import vllm_omni.engine.stage_engine_core_client as client_mod - importlib.reload(client_mod) + # Snapshot the module namespace so a reload cannot strand other test + # modules that captured the pre-reload class objects (isinstance checks + # against the stale class would then fail). ``importlib.reload`` mutates + # the module dict in place; restoring it leaves the runtime exactly as it + # was before this test ran. + saved = dict(client_mod.__dict__) + try: + importlib.reload(client_mod) - assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == ( - "StageEngineCoreClient | DPLBStageEngineCoreClient" - ) + assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == ( + "StageEngineCoreClient | DPLBStageEngineCoreClient" + ) + finally: + client_mod.__dict__.clear() + client_mod.__dict__.update(saved) def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch): diff --git a/tests/engine/test_omni_request.py b/tests/engine/test_omni_request.py index 521c92b94..b1a5abafd 100644 --- a/tests/engine/test_omni_request.py +++ b/tests/engine/test_omni_request.py @@ -13,13 +13,16 @@ substitutable with the base class — including base-style calls that pass import inspect +import msgspec import numpy as np import pytest import torch from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreRequest +from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope from vllm.v1.request import Request -from vllm_omni.engine import PromptEmbedsPayload +from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload from vllm_omni.request import OmniRequest, OmniStreamingUpdate pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -111,3 +114,26 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update(): assert update is not None assert request.model_intermediate_buffer is buffer assert update.model_intermediate_buffer is buffer + + +def test_kv_hints_survive_omni_upgrade_wire_transport_and_scheduler_conversion(): + hints = KvHintsEnvelope("1", "hint-message", [KvHintAction("a", "prefetch", "1", {"blocks": [7]})]) + core = EngineCoreRequest( + request_id="hint-request", + prompt_token_ids=[1, 2], + mm_features=None, + sampling_params=SamplingParams(max_tokens=2), + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + session_id="session", + kv_hints=hints, + ) + upgraded = OmniEngineCoreRequest.from_request(core) + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest) + scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None) + assert decoded.session_id == "session" + assert scheduled.kv_hints == hints + assert scheduled.session_id == "session" diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index eb1bdfd41..0179a221b 100644 --- a/tests/engine/test_omni_request_output.py +++ b/tests/engine/test_omni_request_output.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for OmniRequestOutput class.""" +import numpy as np import pytest from PIL import Image from vllm.outputs import CompletionOutput, RequestOutput @@ -157,6 +158,7 @@ class TestMsgpackRoundTrip: ) source = _make_text_request_output(text="round trip") + source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32) out = OmniRequestOutput.from_stage_output( source, request_id="rt-1", @@ -178,6 +180,7 @@ class TestMsgpackRoundTrip: assert decoded.outputs[0].text == "round trip" assert decoded.prompt == "test prompt" assert decoded.prompt_token_ids == [1, 2, 3] + np.testing.assert_array_equal(decoded.prompt_token_id_logprobs, source.prompt_token_id_logprobs) def test_round_trip_diffusion_output(self): """A diffusion output round-trips without data loss.""" @@ -302,6 +305,16 @@ class TestRequestOutputFieldParity: assert omni.num_cached_tokens is None assert omni.num_cache_creation_tokens is None + def test_fixed_token_scores_survive_stage_output_wrapping(self) -> None: + source = _make_text_request_output() + source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32) + + omni = OmniRequestOutput.from_stage_output(source, stage_id=0) + forwarded = OmniRequestOutput.from_stage_output(omni, stage_id=1) + + np.testing.assert_array_equal(forwarded.prompt_token_id_logprobs, source.prompt_token_id_logprobs) + assert OmniRequestOutput(request_id="unscored").prompt_token_id_logprobs is None + def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None: """Declaring the fields is not enough — they must also be copied. diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index dd1f1efb5..7e1635d19 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -1282,6 +1282,7 @@ async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(metho [ ("reset_mm_cache", {}, True), ("sleep", {}, True), + ("release_kv_cache_memory", {}, True), ("pause_scheduler", {"clear_cache": True}, True), ("pause_scheduler", {"clear_cache": False}, False), ("reset_prefix_cache", {}, False), diff --git a/tests/engine/test_output_processor_mrv2_text.py b/tests/engine/test_output_processor_mrv2_text.py index b9dbecb1b..acfbbe5aa 100644 --- a/tests/engine/test_output_processor_mrv2_text.py +++ b/tests/engine/test_output_processor_mrv2_text.py @@ -7,8 +7,9 @@ from unittest.mock import MagicMock import pytest import torch -from vllm.sampling_params import RequestOutputKind -from vllm.v1.engine import FinishReason +from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm.v1.engine import EngineCoreRequest, FinishReason +from vllm.v1.engine.logprobs import LogprobsProcessor from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.output_modality import OutputModality @@ -16,6 +17,7 @@ from vllm_omni.engine.output_processor import ( MultimodalOutputProcessor, OmniRequestState, ) +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -109,6 +111,37 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload(): assert not state.mm_accumulated.is_empty +def test_fixed_token_scores_survive_output_processing_and_stage_wrapping(): + processor, state = _make_processor(OutputModality.TEXT, _Detokenizer()) + request = EngineCoreRequest( + request_id="r", + prompt_token_ids=[1, 2, 3], + mm_features=None, + sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]), + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + ) + state.logprobs_processor = LogprobsProcessor.from_new_request(None, request) + scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]]) + processed = processor.process_outputs( + [ + OmniEngineCoreOutput( + request_id="r", + new_token_ids=[42], + prompt_token_id_logprobs=scores, + finish_reason=FinishReason.STOP, + ) + ] + ) + + (raw_output,) = processed.request_outputs + final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0) + torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores) + + @pytest.mark.parametrize("streaming", [False, True]) def test_generation_stage_mm_delivered_only_with_terminal(streaming): # Non-terminal chunks accumulate; the terminal output carries the audio. diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 3e8f0ef87..9b8b70586 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -97,7 +97,6 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( "enable_multithread_weight_load", "env", "has_sampling_extra_args", - "log_level", "log_stats", "model_arch", "model_subdir", diff --git a/tests/engine/test_stage_engine_core_client.py b/tests/engine/test_stage_engine_core_client.py index 36c95cde7..37dd07bc7 100644 --- a/tests/engine/test_stage_engine_core_client.py +++ b/tests/engine/test_stage_engine_core_client.py @@ -5,11 +5,19 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import patch import pytest +from vllm import SamplingParams +from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.stage_engine_core_client import ( + DPLBStageEngineCoreClient, + StageEngineCoreClient, + StageEngineCoreClientBase, +) +from vllm_omni.engine.stage_init_utils import StageMetadata pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -30,3 +38,46 @@ def test_check_health_raises_when_resources_engine_dead(): client = _make_client(engine_dead=True) with pytest.raises(EngineDeadError, match="engine core is dead"): client.check_health() + + +@pytest.mark.parametrize("dp_size, external_lb", [(1, False), (2, False), (2, True)]) +@pytest.mark.parametrize("renderer", [None, object()]) +def test_stage_client_forwards_renderer_to_the_actual_upstream_constructor(dp_size, external_lb, renderer): + config = SimpleNamespace( + model_config=None, + parallel_config=SimpleNamespace(data_parallel_size=dp_size, data_parallel_external_lb=external_lb), + ) + metadata = StageMetadata( + stage_id=0, + stage_type="llm", + engine_output_type="text", + is_comprehension=True, + requires_multimodal_data=False, + engine_input_source=[], + final_output=True, + final_output_type="text", + default_sampling_params=SamplingParams(), + custom_process_input_func=None, + model_stage=None, + runtime_cfg=None, + ) + internal_lb = dp_size > 1 and not external_lb + parent = DPLBAsyncMPClient if internal_lb else AsyncMPClient + # Keep Omni construction/factory routing real; intercept only upstream's + # subprocess/socket boundary, using its current signature via autospec. + with patch.object(parent, "__init__", autospec=True, return_value=None) as init: + addresses = {"input_address": "tcp://192.0.2.1:1"} + client = StageEngineCoreClientBase.make_async_mp_client( + config, object, metadata=metadata, client_addresses=addresses, renderer=renderer + ) + assert isinstance(client, DPLBStageEngineCoreClient if internal_lb else StageEngineCoreClient) + init.assert_called_once_with( + client, + config, + object, + log_stats=False, + client_addresses=addresses, + client_count=1, + client_index=0, + renderer=renderer, + ) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index d15df2bdd..aa7bbcb39 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -83,6 +83,27 @@ def test_collective_rpc_control_method_reraises_worker_error(): asyncio.run(run()) +@pytest.mark.cpu +def test_kv_memory_release_uses_engine_core_and_propagates_pause_errors(): + async def run() -> None: + release = AsyncMock(return_value=None) + pool, client = _make_pool(release_kv_cache_memory_async=release) + await pool.collective_rpc(0, "release_kv_cache_memory", timeout=2.0) + release.assert_awaited_once_with() + client.collective_rpc_async.assert_not_awaited() + + release.side_effect = RuntimeError("requires a completed pause first") + with pytest.raises(RuntimeError, match="completed pause"): + await pool.collective_rpc(0, "release_kv_cache_memory") + + del client.release_kv_cache_memory_async + result = await pool.collective_rpc(0, "release_kv_cache_memory") + assert result["supported"] is False + client.collective_rpc_async.assert_not_awaited() + + asyncio.run(run()) + + @pytest.mark.cpu def test_collective_rpc_reset_caches_use_control_helpers(): async def run() -> None: diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 1bcc6e08c..829e30145 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -10,6 +10,7 @@ import pytest from vllm.engine.protocol import StreamingInput from vllm.lora.request import LoRARequest from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm.v1.kv_hints import KvHintsEnvelope from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path @@ -138,6 +139,39 @@ def test_generate_forwards_lora_request_to_engine(): asyncio.run(run()) +@pytest.mark.cpu +@pytest.mark.parametrize("streaming", [None, "chunk", "empty"]) +def test_generate_forwards_kv_hints_on_initial_stage_submission(streaming): + async def run(): + submissions = [] + + async def add(**kwargs): + submissions.append(kwargs) + + omni = get_async_omni_instance(fake_add_request=add) + omni.engine.add_streaming_update_async = _noop + omni.engine.stage_configs = [SimpleNamespace(is_comprehension=True, stage_type="llm")] + omni.engine.stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))] + hints = KvHintsEnvelope("1", "message", []) + params = SamplingParams(max_tokens=2, output_kind=RequestOutputKind.DELTA) + + async def chunks(): + if streaming == "chunk": + yield StreamingInput(prompt={"prompt_token_ids": [1]}) + + async for _ in omni.generate( + prompt=chunks() if streaming else {"prompt_token_ids": [1]}, + request_id="hints", + sampling_params_list=[params], + kv_hints=hints, + ): + pass + assert len(submissions) == 1 + assert submissions[0]["kv_hints"] is hints + + asyncio.run(run()) + + @pytest.mark.cpu @pytest.mark.parametrize( "req_ids,cancel_prefix,expected_cancel_count", diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 7da8f1396..bc4665dcc 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -194,6 +194,56 @@ def test_reset_encoder_cache_forwards_to_ar_stages(): asyncio.run(run()) +@pytest.mark.cpu +def test_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) + with pytest.raises(RuntimeError, match="completed pause"): + await omni.release_kv_cache_memory() + omni.collective_rpc.assert_not_awaited() + + omni._paused = True + await omni.release_kv_cache_memory() + omni._clear_frontend_mm_cache.assert_awaited_once_with() + omni.collective_rpc.assert_awaited_once_with( + method="release_kv_cache_memory", + args=(), + kwargs=None, + stage_ids=[0, 2], + timeout=CACHE_RESET_TIMEOUT_S, + ) + assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}} + assert not omni._level2_sleeping + + omni.collective_rpc.reset_mock() + await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"]) + omni.collective_rpc.assert_awaited_once_with( + method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2] + ) + assert not omni._sleeping_tags + assert omni._paused + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_release_kv_memory_skips_diffusion_and_reports_engine_failures(): + async def run() -> None: + diffusion = _make_omni(stage_types=["diffusion"]) + await diffusion.release_kv_cache_memory() + diffusion.collective_rpc.assert_not_awaited() + assert not diffusion._paused + + omni = _make_omni(stage_types=["llm"]) + omni._paused = True + omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}] + with pytest.raises(RuntimeError, match="all executor memory resident"): + await omni.release_kv_cache_memory() + assert not omni._sleeping_tags + + asyncio.run(run()) + + @pytest.mark.cpu def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages(): async def run() -> None: diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index f6559e0c3..9587d5591 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -4,9 +4,11 @@ """Assertion and response validation helpers for tests.""" import base64 +import hashlib import io import json import math +import os import tempfile import threading import wave @@ -999,6 +1001,31 @@ def _assert_transcript_matches( ) +def _retain_failed_speech_audio( + audio_bytes: bytes | None, transcript: str, expected_text: Any, error: AssertionError +) -> None: + """Retain the exact ASR input when a CI speech-content assertion fails.""" + directory = os.environ.get("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR") + if not directory or not audio_bytes: + return + stem = f"failed_speech_{hashlib.sha256(audio_bytes).hexdigest()[:16]}" + root = Path(directory) + try: + root.mkdir(parents=True, exist_ok=True) + (root / f"{stem}.wav").write_bytes(audio_bytes) + (root / f"{stem}.json").write_text( + json.dumps( + {"transcript": transcript, "expected_text": str(expected_text), "assertion": str(error)}, + ensure_ascii=False, + indent=2, + ) + ) + print(f"Retained failed speech ASR input: {root / f'{stem}.wav'}") + except OSError as exc: + # Diagnostic I/O must never replace or suppress the content failure. + print(f"Could not retain failed speech ASR input: {exc}") + + def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None: """Validate speech API results from :class:`~tests.helpers.client.OmniResponse`. @@ -1063,14 +1090,21 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], if expected_text: print(f"audio content is: {transcript}") print(f"input text is: {expected_text}") - _assert_transcript_matches( - transcript, - _speech_audio_for_transcription(getattr(response, "audio_bytes", None), request_config), - expected_text, - threshold=0.9, - escalation_model=request_config.get("transcript_escalation_model"), - language=request_config.get("transcript_language"), + transcription_audio = _speech_audio_for_transcription( + getattr(response, "audio_bytes", None), request_config ) + try: + _assert_transcript_matches( + transcript, + transcription_audio, + expected_text, + threshold=0.9, + escalation_model=request_config.get("transcript_escalation_model"), + language=request_config.get("transcript_language"), + ) + except AssertionError as exc: + _retain_failed_speech_audio(transcription_audio, transcript, expected_text, exc) + raise _assert_preset_voice_gender_from_audio( response.audio_bytes, request_config.get("voice"), diff --git a/tests/helpers/cpu_collective_server.py b/tests/helpers/cpu_collective_server.py new file mode 100644 index 000000000..30cc80431 --- /dev/null +++ b/tests/helpers/cpu_collective_server.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Import preload for the CPU-only collective forkserver. + +This module is imported in the forkserver process, before it starts workers. +Its CUDA visibility does not alter the pytest process or CUDA parity workers. +""" + +import os + +os.environ["CUDA_VISIBLE_DEVICES"] = "" + +import torch # noqa: E402 + +# Retain the expensive worker-module imports across CPU process groups. +from tests.diffusion.distributed import test_comm, test_pipeline_parallel # noqa: E402, F401 + +if torch.cuda.is_initialized(): + raise RuntimeError("CPU collective forkserver must not initialize CUDA") diff --git a/tests/helpers/distributed.py b/tests/helpers/distributed.py new file mode 100644 index 000000000..1cc79753b --- /dev/null +++ b/tests/helpers/distributed.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Process and platform setup for isolated collective test workers.""" + +from collections.abc import Callable +from multiprocessing.context import BaseContext + +import torch + + +def collective_context(device_kind: str) -> BaseContext: + """Reuse imports in a clean CPU forkserver; CUDA workers use spawn.""" + if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods(): + torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"]) + return torch.multiprocessing.get_context("forkserver") + return torch.multiprocessing.get_context("spawn") + + +def start_collective_workers( + worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str +) -> None: + context = collective_context(device_kind) + if context.get_start_method() == "spawn": + torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs) + else: + torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method()) + + +def configure_cpu_collective_worker() -> None: + """Select CPU platforms inside a spawned worker that only uses Gloo. + + Upstream coordinators resolve local ranks through the selected + platform's visible devices. CPU ranks can outnumber visible GPUs, so they + must use the CPU platform even when the test host has CUDA installed. + The worker still creates real process groups and runs real collectives. + """ + import vllm.platforms + from vllm.platforms.cpu import CpuPlatform + + from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state + from vllm_omni.platforms.interface import UnspecifiedOmniPlatform + + class CpuCollectivePlatform(UnspecifiedOmniPlatform): + dist_backend = "gloo" + + @classmethod + def get_device_count(cls) -> int: + # CPU tensors share one logical device regardless of process rank. + return 1 + + @classmethod + def set_device(cls, device: torch.device) -> None: + assert device.type == "cpu" + + @classmethod + def synchronize(cls) -> None: + pass + + vllm.platforms.current_platform = CpuPlatform() + platform = CpuCollectivePlatform() + for module in (comm, group_coordinator, parallel_state): + module.current_omni_platform = platform diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 5fedd17dd..2475431dd 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -659,6 +659,14 @@ class OmniRunner: self.seed = seed self._prompt_len_estimate_cache: dict[str, Any] = {} self.omni: Any = None + self._runner_process = psutil.Process() + self._initial_child_ids: set[tuple[int, float]] = set() + self._owned_engine_processes: dict[int, float] = {} + for child in self._runner_process.children(recursive=True): + try: + self._initial_child_ids.add((child.pid, child.create_time())) + except psutil.Error: + pass try: from vllm_omni.entrypoints.omni import Omni @@ -670,12 +678,14 @@ class OmniRunner: deploy_config=deploy_config, **kwargs, ) + self._remember_engine_processes(startup=True) startup_s = time.perf_counter() - startup_t0 if log_stats: print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True) except BaseException: # ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__`` # when construction fails after worker processes have started. + self._remember_engine_processes(startup=True) self.__exit__(None, None, None) raise @@ -899,18 +909,51 @@ class OmniRunner: def stop_profile(self, stages: list[int] | None = None) -> list[Any]: return self.omni.stop_profile(stages=stages) + def _remember_engine_processes(self, *, startup: bool = False) -> None: + """Track startup workers, then only descendants of those workers. + + Existing children may belong to another runner or a shared transcriber. + Record identities before close() can reparent surviving workers. + """ + try: + if startup: + candidates = self._runner_process.children(recursive=True) + else: + candidates = [] + for pid, created in list(self._owned_engine_processes.items()): + try: + proc = psutil.Process(pid) + if proc.create_time() == created: + candidates.extend(proc.children(recursive=True)) + except psutil.Error: + pass + for proc in candidates: + try: + identity = (proc.pid, proc.create_time()) + if identity in self._initial_child_ids: + continue + if startup and any( + (parent.pid, parent.create_time()) in self._initial_child_ids for parent in proc.parents() + ): + continue + cmdline = " ".join(proc.cmdline()).lower() + if "enginecore" in cmdline or "enginecore" in proc.name().lower(): + self._owned_engine_processes[identity[0]] = identity[1] + except psutil.Error: + pass + except psutil.Error as error: + print(f"Warning: could not inspect runner workers: {error}") + def _cleanup_process(self): try: - keywords = ["enginecore"] matched = [] - for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]): + for pid, created in self._owned_engine_processes.items(): try: - cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else "" - name = proc.name().lower() - if any(k in cmdline for k in keywords) or any(k in name for k in keywords): - print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}") + proc = psutil.Process(pid) + if proc.create_time() == created: + print(f"Found owned vllm process: PID={proc.pid}") matched.append(proc) - except (psutil.NoSuchProcess, psutil.AccessDenied): + except psutil.Error: pass for proc in matched: try: @@ -938,11 +981,14 @@ class OmniRunner: return self def __exit__(self, exc_type, exc_val, exc_tb): - omni = getattr(self, "omni", None) - if omni is not None and hasattr(omni, "close"): - omni.close() - self._cleanup_process() - cleanup_test_environment() + self._remember_engine_processes() + try: + omni = getattr(self, "omni", None) + if omni is not None and hasattr(omni, "close"): + omni.close() + finally: + self._cleanup_process() + cleanup_test_environment() # --------------------------------------------------------------------------- diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index e2060152f..2fca68d65 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import io +import json import wave +from contextlib import nullcontext from types import SimpleNamespace import numpy as np @@ -249,6 +251,85 @@ def test_pcm_transcript_escalation_uses_same_valid_wav_and_rejects_two_wrong_tra _assert_wav_wraps_pcm(wav_bytes, pcm_bytes, 24_000) +@pytest.mark.parametrize("response_format", ["wav", "pcm"]) +@pytest.mark.parametrize("escalation_model", [None, "large-v3"]) +def test_failed_speech_capture_retains_exact_asr_input_and_original_failure( + monkeypatch, tmp_path, response_format, escalation_model +): + pcm_bytes = _pcm_sine() + wav = io.BytesIO() + with wave.open(wav, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(24_000) + writer.writeframes(pcm_bytes) + audio_bytes = wav.getvalue() if response_format == "wav" else pcm_bytes + transcribed = [] + + def transcribe(audio_bytes, **kwargs): + transcribed.append(audio_bytes) + return "this is unrelated speech" + + directory = tmp_path / "failed_audio" + monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory)) + monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", transcribe) + config = { + "input": "please read this requested sentence accurately", + "response_format": response_format, + "transcript_pcm_sample_rate": 24_000, + "transcript_escalation_model": escalation_model, + } + response = OmniResponse(success=True, audio_bytes=audio_bytes, audio_format=f"audio/{response_format}") + with pytest.raises(AssertionError, match="Transcript doesn't match input") as error: + assert_audio_speech_response(response, config, "full_model") + + clips = list(directory.glob("*.wav")) + assert len(clips) == 1 + assert all(clip == clips[0].read_bytes() for clip in transcribed) + _assert_wav_wraps_pcm(clips[0].read_bytes(), pcm_bytes, 24_000) + metadata = json.loads(clips[0].with_suffix(".json").read_text()) + assert metadata["expected_text"] == config["input"] + assert metadata["transcript"] == "this is unrelated speech" + assert metadata["assertion"] == str(error.value) + + +@pytest.mark.parametrize("capture_enabled", [False, True]) +@pytest.mark.parametrize("content_failure", [False, True]) +def test_speech_capture_is_opt_in_and_does_not_write_successes(monkeypatch, tmp_path, capture_enabled, content_failure): + directory = tmp_path / "failed_audio" + if capture_enabled: + monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory)) + else: + monkeypatch.delenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", raising=False) + expected = "please read this requested sentence accurately" + monkeypatch.setattr( + assertions, + "convert_audio_bytes_to_text", + lambda *_args, **_kwargs: "this is unrelated speech" if content_failure else expected, + ) + with pytest.raises(AssertionError, match="Transcript doesn't match input") if content_failure else nullcontext(): + assert_audio_speech_response( + OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"), + {"input": expected, "response_format": "wav"}, + "full_model", + ) + assert directory.exists() is (capture_enabled and content_failure) + assert len(list(tmp_path.rglob("*.wav"))) == int(capture_enabled and content_failure) + + +def test_speech_capture_io_failure_does_not_mask_content_failure(monkeypatch, tmp_path): + directory = tmp_path / "blocked" + directory.write_text("a file cannot be an output directory") + monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory)) + monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "this is unrelated speech") + with pytest.raises(AssertionError, match="Transcript doesn't match input"): + assert_audio_speech_response( + OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"), + {"input": "please read this requested sentence accurately", "response_format": "wav"}, + "full_model", + ) + + @pytest.mark.parametrize( "pcm_bytes", [ diff --git a/tests/helpers/tests/test_runtime_omni_runner.py b/tests/helpers/tests/test_runtime_omni_runner.py index aaada20f2..bcd29493d 100644 --- a/tests/helpers/tests/test_runtime_omni_runner.py +++ b/tests/helpers/tests/test_runtime_omni_runner.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Tests for OmniRunner startup rollback.""" +"""Tests for OmniRunner startup rollback and worker ownership.""" from __future__ import annotations +import subprocess +import sys +from collections.abc import Generator + +import psutil import pytest from tests.helpers.runtime import OmniRunner @@ -38,3 +43,114 @@ def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.Monk assert cleaned.count("process") == 1 # Once at the start of ``__init__``, once from the constructor rollback. assert cleaned.count("env") == 2 + + +@pytest.fixture +def live_workers() -> Generator[list[subprocess.Popen[bytes]], None, None]: + workers: list[subprocess.Popen[bytes]] = [] + try: + yield workers + finally: + for worker in workers: + if worker.poll() is None: + worker.kill() + worker.wait(timeout=10) + if worker.stdout is not None: + worker.stdout.close() + + +def _start_engine_worker( + workers: list[subprocess.Popen[bytes]], *, ignore_terminate: bool = False +) -> subprocess.Popen[bytes]: + script = "import signal, time; " + if ignore_terminate: + script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); " + script += "print('ready', flush=True); time.sleep(90)" + worker = subprocess.Popen([sys.executable, "-c", script, "enginecore"], stdout=subprocess.PIPE) + workers.append(worker) + assert worker.stdout is not None + assert worker.stdout.readline() == b"ready\n" + return worker + + +@pytest.mark.parametrize("close_raises", [False, True]) +def test_runner_cleanup_preserves_unrelated_engine_workers( + monkeypatch: pytest.MonkeyPatch, + live_workers: list[subprocess.Popen[bytes]], + close_raises: bool, +) -> None: + existing = _start_engine_worker(live_workers) + owned: list[subprocess.Popen[bytes]] = [] + late_unrelated: list[subprocess.Popen[bytes]] = [] + monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None) + + class _OwnedOmni: + def __init__(self, *args: object, **kwargs: object) -> None: + owned.append(_start_engine_worker(live_workers)) + + def close(self) -> None: + late_unrelated.append(_start_engine_worker(live_workers)) + if close_raises: + raise RuntimeError("close failed") + + monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _OwnedOmni) + if close_raises: + with pytest.raises(RuntimeError, match="close failed"), OmniRunner("fake-model"): + pass + else: + with OmniRunner("fake-model"): + pass + + assert owned[0].poll() is not None + assert existing.poll() is None + assert late_unrelated[0].poll() is None + + +def test_runner_startup_failure_cleans_only_new_workers( + monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]] +) -> None: + unrelated = _start_engine_worker(live_workers) + owned: list[subprocess.Popen[bytes]] = [] + monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None) + + class _FailingOmni: + def __init__(self, *args: object, **kwargs: object) -> None: + owned.append(_start_engine_worker(live_workers)) + raise RuntimeError("startup failed") + + monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _FailingOmni) + with pytest.raises(RuntimeError, match="startup failed"): + OmniRunner("fake-model") + + assert owned[0].poll() is not None + assert unrelated.poll() is None + + +def test_runner_cleanup_rejects_stale_process_identity(live_workers: list[subprocess.Popen[bytes]]) -> None: + unrelated = _start_engine_worker(live_workers) + runner = object.__new__(OmniRunner) + runner._owned_engine_processes = {unrelated.pid: psutil.Process(unrelated.pid).create_time() + 1} + + runner._cleanup_process() + + assert unrelated.poll() is None + + +def test_runner_cleanup_kills_owned_worker_ignoring_terminate( + monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]] +) -> None: + owned: list[subprocess.Popen[bytes]] = [] + monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None) + + class _StubbornOmni: + def __init__(self, *args: object, **kwargs: object) -> None: + owned.append(_start_engine_worker(live_workers, ignore_terminate=True)) + + def close(self) -> None: + pass + + monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _StubbornOmni) + with OmniRunner("fake-model"): + assert owned[0].poll() is None + + assert owned[0].poll() is not None diff --git a/tests/inputs/test_mm_processor.py b/tests/inputs/test_mm_processor.py index 813d7f436..de2de98fb 100644 --- a/tests/inputs/test_mm_processor.py +++ b/tests/inputs/test_mm_processor.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from importlib import import_module from types import SimpleNamespace from unittest.mock import Mock import pytest from transformers import BatchFeature +from vllm.config.multimodal import MultiModalDummyOptions +from vllm.multimodal.parse import AudioProcessorItems, ImageProcessorItems, MultiModalDataItems +from vllm.multimodal.processing import BaseMultiModalProcessor -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -24,10 +28,14 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, processor = object.__new__(Processor) processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy")) - processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"})) processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"})) - items = Mock() - items.get_all_counts.return_value = counts + + class AudioItems(AudioProcessorItems): + def get_passthrough_data(self): + return {"embedding": "preserve"} + + items = MultiModalDataItems({"audio": AudioItems(["wave"]), "image": ImageProcessorItems([])} if counts else {}) + items.select = Mock(wraps=items.select) kwargs = {"sampling_rate": 16000} if original_text is not None: kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text @@ -36,9 +44,83 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, items.select.assert_called_once_with({k for k, v in counts.items() if v > 0}) processor._call_hf_processor.assert_called_once_with( "dummy" if original_text is None else original_text, - {"audios": ["wave"]}, + {"audios": ["wave"]} if counts else {}, {"sampling_rate": 16000}, {}, ) - assert dict(result) == {"features": "processed", "embedding": "preserve"} + assert dict(result) == ({"features": "processed", "embedding": "preserve"} if counts else {"features": "processed"}) assert kwargs == before + + +@pytest.mark.parametrize( + "module_name,builder_name,processor_name,expected_kwargs", + [ + ( + "cosyvoice3", + "CosyVoice3DummyInputsBuilder", + "CosyVoice3MultiModalProcessor", + {"prompt_text": "Testing my voices. Why should I not?"}, + ), + ( + "glm_tts", + "GLMTTSDummyInputsBuilder", + "GLMTTSMultiModalProcessor", + {"prompt_text": "This is the reference voice."}, + ), + ( + "omnivoice", + "OmniVoiceDummyInputsBuilder", + "OmniVoiceMultiModalProcessor", + {"ref_text": "Testing voice cloning."}, + ), + ("mimo_audio", "MiMoAudioLLMDummyInputsBuilder", "MiMoAudioLLMMultiModalProcessor", {}), + ], +) +def test_dummy_budget_preserves_model_conditioning(module_name, builder_name, processor_name, expected_kwargs): + """The v0.31 budget path must reach each model's custom builder hook.""" + module = import_module(f"vllm_omni.model_executor.models.{module_name}.{module_name}") + builder = object.__new__(getattr(module, builder_name)) + assert isinstance(builder, OmniDummyInputsBuilder) + items = MultiModalDataItems({}) + builder.info = SimpleNamespace(ctx=SimpleNamespace(tokenizer=None), parse_mm_data=Mock(return_value=items)) + builder.get_dummy_mm_data = Mock(return_value={}) + processor = object.__new__(getattr(module, processor_name)) + processor.dummy_inputs = builder + + inputs = processor.get_dummy_inputs(128, {"audio": 1}, MultiModalDummyOptions()) + + assert inputs.mm_data_items is items + assert inputs.hf_processor_mm_kwargs == expected_kwargs + assert inputs.prompt == ("<|empty|>" if module_name == "mimo_audio" else []) + + +def test_dummy_budget_keeps_upstream_path_for_standard_builders(monkeypatch): + class Processor(OmniMultiModalProcessor): + def _get_mm_fields_config(self, *args): + return {} + + def _get_prompt_updates(self, *args): + return [] + + processor = object.__new__(Processor) + processor.dummy_inputs = SimpleNamespace() + expected = object() + upstream = Mock(return_value=expected) + monkeypatch.setattr(BaseMultiModalProcessor, "get_dummy_inputs", upstream) + options = MultiModalDummyOptions() + + assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected + upstream.assert_called_once_with(128, {"audio": 1}, options) + + +def test_voxtral_dummy_budget_uses_audio_token_prompt(): + from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import VoxtralTTSMultiModalProcessor + + processor = object.__new__(VoxtralTTSMultiModalProcessor) + expected = SimpleNamespace(prompt=[1, 25, 24, 35]) + builder = SimpleNamespace(get_dummy_processor_inputs=Mock(return_value=expected)) + processor.dummy_inputs = builder + options = MultiModalDummyOptions() + + assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected + builder.get_dummy_processor_inputs.assert_called_once_with(128, {"audio": 1}, options) diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index 3fb48a7cc..c61920ba9 100644 --- a/tests/inputs/test_preprocess.py +++ b/tests/inputs/test_preprocess.py @@ -8,7 +8,9 @@ from unittest.mock import AsyncMock, Mock import pytest import torch -from vllm.inputs import tokens_input +from vllm.exceptions import VLLMValidationError +from vllm.inputs import mm_input, tokens_input +from vllm.multimodal.inputs import PlaceholderRange from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.renderers.params import TokenizeParams @@ -115,6 +117,35 @@ def test_omni_renderer_is_a_real_subclass_created_once(): assert not hasattr(instance, "_renderer") +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("no_media_kwargs", [False, True]) +@pytest.mark.parametrize("limit", [2, 4]) +def test_renderer_applies_expanded_prompt_bound_without_splitting_media(renderer, async_mode, no_media_kwargs, limit): + renderer._process_multimodal.return_value = mm_input( + list(range(6)), {}, {"image": ["hash"]}, {"image": [PlaceholderRange(offset=1, length=2)]} + ) + prompt = {"prompt_token_ids": [1, 2, 3], "additional_information": {"speaker": 1}} + if no_media_kwargs: + prompt["mm_processor_kwargs"] = {"target_h": 512} + else: + prompt["multi_modal_data"] = {"image": "image"} + params = TokenizeParams(max_total_tokens=128, truncate_prompt_tokens=limit, truncation_side="right") + + def render(): + if async_mode: + return asyncio.run(renderer.render_cmpl_async([prompt], params)) + return renderer.render_cmpl([prompt], params) + + if limit == 2: + with pytest.raises(VLLMValidationError, match="would split or drop"): + render() + else: + (result,) = render() + assert result["prompt_token_ids"] == [0, 1, 2, 3] + assert result["mm_placeholders"]["image"] == [PlaceholderRange(offset=1, length=2)] + assert result["additional_information"] == {"speaker": 1} + + def test_concrete_renderer_overrides_are_not_suppressed(): """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides.""" diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index 515aaeb8d..ecca06e5b 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -426,6 +426,8 @@ def _make_scheduler_shell() -> OmniARScheduler: obj.requests = {} obj.running = [] obj.waiting = [] + obj.kv_holding_waiting = [] + obj.deferred_waiting = set() obj._kv_wait_start_ts = {} obj._omni_kv_config = None return obj diff --git a/tests/metrics/test_stat_logger.py b/tests/metrics/test_stat_logger.py index 97a18911d..86c43a500 100644 --- a/tests/metrics/test_stat_logger.py +++ b/tests/metrics/test_stat_logger.py @@ -477,6 +477,7 @@ def stats_config(): observability_config=ObservabilityConfig(kv_cache_metrics=True), speculative_config=Mock(spec=SpeculativeConfig, num_speculative_tokens=3), kv_transfer_config=None, + ec_transfer_config=None, lora_config=None, ) yield config diff --git a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py index 2d529cf26..f154e7ee4 100644 --- a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py +++ b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py @@ -7,7 +7,11 @@ from types import SimpleNamespace import numpy as np import pytest import torch +from vllm.config import VllmConfig from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.bad_words import BadWordsState +from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState +from vllm.v1.worker.gpu.sample.penalties import PenaltiesState from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.states import RequestState @@ -27,11 +31,11 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, pytest.skip("CUDA required") device = torch.device("cuda") reqs = RequestState(5, 32, 16, 0, 8, device) - sampler = Sampler(5, 8, device, reqs) + sampler = Sampler(VllmConfig(), 5, 8, device, reqs) for i, length in enumerate([4, 1, 3]): reqs.add_request(str(i), length, [2] * length, 0, 8) slot = reqs.req_id_to_index[str(i)] - sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed)) + sampler.add_request(slot, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed)) reqs.apply_staged_writes() sampler.apply_staged_writes() state = _state(MossLocalModelState, device) @@ -39,6 +43,7 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, state.model._batch_state = None batch = _batch(device, [2, 4, 3], [1, 2, 1]) batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32) + batch.seq_lens_cpu_upper_bound = torch.tensor([5, 2, 2], dtype=torch.int32) batch.cu_num_logits_np = np.arange(4, dtype=np.int32) batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device) batch.expanded_idx_mapping = batch.idx_mapping @@ -73,10 +78,10 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint): pytest.skip("CUDA required") device = torch.device("cuda") reqs = RequestState(5, 32, 16, 0, 8, device) - sampler = Sampler(5, 8, device, reqs) + sampler = Sampler(VllmConfig(), 5, 8, device, reqs) reqs.add_request("constrained", 1, [2], 0, 8) slot = reqs.req_id_to_index["constrained"] - sampler.add_request(slot, 1, SamplingParams(**constraint)) + sampler.add_request(slot, SamplingParams(**constraint)) state = _state(MossLocalModelState, device) state._direct_tokens = True # Use the actual upstream request registration, including its combined @@ -85,16 +90,22 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint): @pytest.mark.cpu -@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"]) +@pytest.mark.parametrize( + "feature", ["logprobs", "bias", "penalty", "bad_words", "custom", "thinking", "trace", "mask", "nans"] +) def test_distribution_features_fall_back_before_using_gpu(feature): sampler = object.__new__(Sampler) sampler.compute_nans = feature == "nans" sampler.return_sampling_mask = feature == "mask" sampler.trace_replay_state = object() if feature == "trace" else None sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None - sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"])) - sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"])) - sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"]))) + bias = object.__new__(LogitBiasState) + bias.use_logit_bias = np.array([feature == "bias"]) + penalties = object.__new__(PenaltiesState) + penalties.use_penalty = np.array([feature == "penalty"]) + bad_words = object.__new__(BadWordsState) + bad_words.num_bad_words = SimpleNamespace(np=np.array([feature == "bad_words"])) + sampler.logits_processors = [bias, penalties, bad_words] + ([object()] if feature == "custom" else []) sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"])) state = object.__new__(MossLocalModelState) state._direct_tokens = True diff --git a/tests/model_executor/models/moss_tts/test_local_eager_mtp.py b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py index cf3893cbd..08828f375 100644 --- a/tests/model_executor/models/moss_tts/test_local_eager_mtp.py +++ b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py @@ -132,10 +132,10 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc hidden_states=torch.zeros(2, 4), finished_req_ids=set(), ec_connector_output=None, - routed_experts=None, ) runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False + runner.aux_output_connector = None runner.model_config = SimpleNamespace(async_chunk=False) runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text")) runner.model_state, runner.model = state, model @@ -154,7 +154,10 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc torch.zeros(1), ) ) - runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={})) + runner.prompt_logprobs_worker = SimpleNamespace( + compute_prompt_logprobs=mocker.Mock(return_value={}), + compute_prompt_token_id_logprobs=mocker.Mock(return_value={}), + ) runner.postprocess_sampled = mocker.Mock() runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None)) mock_out = SimpleNamespace(copy_event=None) diff --git a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py index 9c6533c76..ab661b11a 100644 --- a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py +++ b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py @@ -42,10 +42,26 @@ def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProce return processor +class _FakeProcessorItem: + def __init__(self, processor_data: dict[str, object]) -> None: + self._processor_data = processor_data + + def get_processor_data(self) -> dict[str, object]: + return dict(self._processor_data) + + def get_passthrough_data(self) -> dict[str, object]: + return {} + + def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace: + class _Selected: + @staticmethod + def values(): + return [_FakeProcessorItem(mm_data)] if mm_data else [] + return SimpleNamespace( get_all_counts=lambda: counts, - select=lambda _keys: mm_data, + select=lambda _keys: _Selected(), ) @@ -75,9 +91,8 @@ def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> Non def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None: processor = _make_processor(monkeypatch) audio = np.zeros(2400, dtype=np.float32) - monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {})) - out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"}) + out = processor._apply_hf_processor_main(_items({"audio": 1}, {"audios": [audio]}), {"ref_text": "reference"}) assert tuple(out["ref_audio_tokens"].shape) == (8, 4) assert list(out["ref_audio_len"]) == [4] @@ -89,7 +104,6 @@ def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.Monke def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: """Text-only requests and fully cached audio both reach this path with empty items.""" processor = _make_processor(monkeypatch) - monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {})) out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"}) @@ -99,4 +113,6 @@ def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.Monke def test_legacy_call_hf_processor_is_gone() -> None: - assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor") + # Upstream reintroduced a 2-arg `_call_hf_processor` on the base class; the + # omni-specific 4-arg hook must not be defined on the omni class itself. + assert "_call_hf_processor" not in OmniVoiceMultiModalProcessor.__dict__ diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index 699acc951..110776488 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -262,6 +262,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): "audio": 1, }[modality] mm_data_items.__getitem__ = MagicMock(side_effect=lambda modality: [f"{modality}-0"]) + mm_data_items.values = MagicMock(return_value=[]) inputs = MagicMock() inputs.mm_data_items = mm_data_items @@ -275,6 +276,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): cache = MagicMock() # Video miss, audio hit under pair-aware keys — expansion must force audio. cache.is_cached.side_effect = lambda keys: [False] if keys[0].startswith("v") else [True] + inputs.cache = cache captured = {} @@ -309,6 +311,8 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): fake_self._get_mm_fields_config = MagicMock(return_value={}) fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]}) fake_self._merge_mm_kwargs = capture_merge + fake_self._postprocess_prompt = lambda prompt: prompt + inputs.prompt = [0, 1, 2] with patch( "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs", @@ -549,24 +553,6 @@ def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count(): ) -def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio(): - captured_counts = {} - fake_self = SimpleNamespace( - dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"), - _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False), - ) - mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1}) - - Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only( - fake_self, - mm_items, - {"use_audio_in_video": [True, False]}, - {}, - ) - - assert captured_counts == {"video": 2, "audio": 0} - - def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits(): fake_self = _fake_processor() update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID])) @@ -661,7 +647,12 @@ def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): Qwen3OmniMoeThinkerMultiModalProcessor, Qwen2_5OmniThinkerMultiModalProcessor, ) - assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data + # Keep Qwen3's native HF preprocessing while inheriting Omni pair handling. + assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_inputs is UpstreamQwen3Processor._get_hf_mm_inputs + assert ( + Qwen3OmniMoeThinkerMultiModalProcessor._get_audio_in_video_pairs + is Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs + ) mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__ assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor) diff --git a/tests/model_executor/models/test_omni_processing.py b/tests/model_executor/models/test_omni_processing.py index eaa798fba..f1493ad5e 100644 --- a/tests/model_executor/models/test_omni_processing.py +++ b/tests/model_executor/models/test_omni_processing.py @@ -79,31 +79,16 @@ def get_text_token_prompts( """Return ``(text_prompt, token_prompt)`` for the given multimodal data.""" dummy_inputs = processor.dummy_inputs tokenizer: TokenizerLike = processor.info.get_tokenizer() - model_config = processor.info.ctx.model_config parsed_data = processor.info.parse_mm_data(mm_data) mm_counts = {k: len(vs) for k, vs in parsed_data.items()} - inputs = dummy_inputs.get_dummy_processor_inputs( - model_config.max_model_len, - mm_counts, - mm_options={}, + text_prompt = dummy_inputs.get_dummy_text(mm_counts) + token_prompt = tokenizer.encode( + text_prompt, + **processor.info.get_default_tok_params().get_encode_kwargs(), ) - text_prompt: str | None - token_prompt: list[int] - if isinstance(inputs.prompt, list): - text_prompt = None - token_prompt = inputs.prompt - elif isinstance(inputs.prompt, str): - text_prompt = inputs.prompt - token_prompt = tokenizer.encode( - text_prompt, - **processor.info.get_default_tok_params().get_encode_kwargs(), - ) - else: - raise TypeError(type(inputs.prompt)) - return text_prompt, token_prompt @@ -199,8 +184,8 @@ def _test_processing_correctness( modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items() } - baseline_processor = factories.build_processor(ctx, cache=None) - cached_processor = factories.build_processor(ctx, cache=cache) + baseline_processor = factories.build_processor(ctx) + cached_processor = factories.build_processor(ctx) rng = np.random.RandomState(0) @@ -238,6 +223,7 @@ def _test_processing_correctness( baseline_processor, cached_processor, batch_idx, + cache, ) @@ -247,6 +233,7 @@ def _test_processing_correctness_one( baseline_processor: BaseMultiModalProcessor, cached_processor: BaseMultiModalProcessor, batch_idx: int, + cache: MultiModalProcessorOnlyCache, ): model_type = model_config.hf_config.model_type @@ -258,12 +245,14 @@ def _test_processing_correctness_one( token_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=None, ) cached_tokenized_result = cached_processor( token_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) _assert_inputs_equal( @@ -278,11 +267,13 @@ def _test_processing_correctness_one( text_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=None, ) cached_text_result = cached_processor( text_prompt, mm_items=mm_items, hf_processor_mm_kwargs={}, + cache=cache, ) _assert_inputs_equal( diff --git a/tests/model_executor/test_cfg_pairing.py b/tests/model_executor/test_cfg_pairing.py index c00313a2a..6d44e70dc 100644 --- a/tests/model_executor/test_cfg_pairing.py +++ b/tests/model_executor/test_cfg_pairing.py @@ -66,6 +66,8 @@ def _make_scheduler_cls() -> type: class FakeScheduler: def __init__(self) -> None: self.waiting = _FakeWaitingQueue() + self.kv_holding_waiting = _FakeWaitingQueue() + self.deferred_waiting: set[_FakeRequest] = set() self.running: list[_FakeRequest] = [] self.requests: dict[str, _FakeRequest] = {} self.finished: list[set[str]] = [] @@ -196,6 +198,44 @@ def test_lone_member_is_held_then_released(scheduler: Any) -> None: assert scheduler.running == [] +def test_lone_kv_holding_member_is_held(scheduler: Any) -> None: + request = _FakeRequest("req-1", {"cfg_role": "cond"}) + scheduler.add_request(request) + scheduler.waiting.remove_requests([request]) + scheduler.kv_holding_waiting.append(request) + + held = cfg_pairing._hold_incomplete_pairs(scheduler) + + assert list(scheduler.kv_holding_waiting) == [] + cfg_pairing._release_held(held) + assert list(scheduler.kv_holding_waiting) == [request] + + +def test_deferred_kv_partner_holds_the_fresh_member(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + scheduler.waiting.remove_requests([uncond]) + scheduler.kv_holding_waiting.append(uncond) + scheduler.deferred_waiting.add(uncond) + + held = cfg_pairing._hold_incomplete_pairs(scheduler) + + assert list(scheduler.waiting) == [] + assert list(scheduler.kv_holding_waiting) == [uncond] + cfg_pairing._release_held(held) + assert list(scheduler.waiting) == [cond] + + +def test_kv_holding_queue_keeps_pair_partners_adjacent(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + other = _FakeRequest("other") + scheduler.kv_holding_waiting.extend([cond, other, uncond]) + scheduler.waiting.clear() + + cfg_pairing._reorder_waiting_for_cfg(scheduler) + + assert list(scheduler.kv_holding_waiting) == [cond, uncond, other] + + def test_complete_pair_is_admitted_together(scheduler: Any) -> None: scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) scheduler.schedule() diff --git a/tests/platforms/npu/test_aux_output.py b/tests/platforms/npu/test_aux_output.py new file mode 100644 index 000000000..8a4362d6d --- /dev/null +++ b/tests/platforms/npu/test_aux_output.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU contract coverage for NPU routing capture and the native data plane.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.config import AuxOutputConfig +from vllm.distributed.aux_output_connector import worker +from vllm.distributed.aux_output_connector.connector import ( + AuxOutputConnectorMetadata, + AuxOutputSchedulerConnector, + PackedBlockHashes, +) +from vllm.model_executor.layers.fused_moe import routed_experts_capturer +from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Router: + def _select_experts(self, logical_ids): + return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids + + +class _AscendMoE(torch.nn.Module): + __module__ = "vllm_ascend.ops.fused_moe" + + def __init__(self): + super().__init__() + self.layer_id = 0 + self.router = _Router() + + +class _Runner(NPUAuxOutputMixin): + def _sync_device(self): + self.synchronized += 1 + + +@pytest.fixture +def runner(monkeypatch): + # Exercise the real capturer, block store, worker and scheduler connectors + # with CPU tensors; importing vllm-ascend or owning NPU hardware is unnecessary. + monkeypatch.setattr(routed_experts_capturer, "current_platform", SimpleNamespace(device_type="cpu")) + monkeypatch.setattr(routed_experts_capturer, "get_forward_context", lambda: SimpleNamespace(dp_metadata=None)) + monkeypatch.setattr(worker, "get_tp_group", lambda: SimpleNamespace(is_first_rank=True)) + runner = _Runner() + runner.model = torch.nn.Sequential(_AscendMoE()) + runner.vllm_config = SimpleNamespace( + aux_output_config=AuxOutputConfig(enable_return_routed_experts=True), + model_config=SimpleNamespace( + get_total_num_hidden_layers=lambda: 1, + get_num_experts=lambda: 16, + get_num_experts_per_tok=lambda: 2, + hf_text_config=SimpleNamespace(model_type="test_moe"), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16, max_num_seqs=4), + parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1, decode_context_parallel_size=1), + cache_config=SimpleNamespace(block_size=4, prefix_match_unit=None), + max_concurrent_batches=2, + ) + runner.input_batch = SimpleNamespace(req_ids=["r"], num_computed_tokens_cpu=np.array([0])) + runner.query_start_loc = SimpleNamespace(np=np.array([0, 4])) + runner.synchronized = 0 + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + try: + yield runner + finally: + runner._close_omni_aux_output() + + +def _step(runner, *, emit_start=0, hashes=None, finished=()): + scheduler_output = SimpleNamespace( + aux_output_connector_metadata=AuxOutputConnectorMetadata( + generation=0, + requests={} if finished else {"r": emit_start}, + block_hashes=hashes or {}, + finished_requests=finished, + ), + scheduled_spec_decode_tokens={}, + ) + runner._begin_omni_aux_output_step(scheduler_output) + return scheduler_output + + +def _capture(runner, ids): + logical_ids = torch.tensor(ids, dtype=torch.int32) + weights, returned = runner.model[0].router._select_experts(logical_ids) + assert torch.equal(returned, logical_ids) + assert weights.shape == logical_ids.shape + return logical_ids.numpy()[:, None, :].astype(np.uint8) + + +def test_routing_reaches_scheduler_and_owns_legacy_batch_snapshot(runner): + step = _step(runner) + expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + # Simulate bookkeeping and the next batch reusing the runner's buffers. + runner.input_batch.num_computed_tokens_cpu[:] = 99 + runner.query_start_loc.np[:] = 99 + _capture(runner, [[7, 8]] * 4) + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]])) + request = SimpleNamespace(request_id="r", num_tokens=5) + rows = AuxOutputSchedulerConnector().take_output(request, output) + np.testing.assert_array_equal(rows, expected) + assert output["r"].token_start == 0 + assert runner.synchronized == 1 + + +def test_speculative_rejection_drops_unaccepted_routing_suffix(runner): + step = _step(runner) + step.scheduled_spec_decode_tokens = {"r": [7, 8, 9]} + expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5, 6, -1, -1]])) + np.testing.assert_array_equal(output["r"].rows, expected[:2]) + + +def test_partial_prefill_accumulates_rows_and_no_forward_step_tears_down(runner): + step = _step(runner) + prefix = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + assert runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]), [0]) == {} + + # Both completed four-token blocks need their scheduler-provided hashes. + step = _step(runner, hashes={"r": PackedBlockHashes(b"h0h1", 2)}) + runner.input_batch.num_computed_tokens_cpu[:] = 4 + suffix = _capture(runner, [[4, 5], [5, 6], [6, 7], [7, 8]]) + pending = runner._prepare_omni_aux_output() + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]])) + np.testing.assert_array_equal(output["r"].rows, np.concatenate((prefix, suffix))) + _step(runner, finished=("r",)) + assert not runner.aux_output_connector._requests + store = runner.aux_output_connector._store + runner._close_omni_aux_output() + assert not store._thread.is_alive() + assert runner.model[0].capture_fn is None + runner._close_omni_aux_output() # shutdown is idempotent + + +def test_reinitialization_does_not_wrap_router_twice(runner): + selection = runner.model[0].router._select_experts + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + assert runner.model[0].router._select_experts is selection + step = _step(runner) + expected = _capture(runner, [[0, 1]] * 4) + output = runner._finish_omni_aux_output(runner._prepare_omni_aux_output(), step, torch.tensor([[5]])) + np.testing.assert_array_equal(output["r"].rows, expected) + + +def test_disabled_capture_preserves_normal_runner_path(runner): + runner._close_omni_aux_output() + runner.vllm_config.aux_output_config.enable_return_routed_experts = False + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + runner._begin_omni_aux_output_step(SimpleNamespace()) + assert runner._prepare_omni_aux_output() is None + assert runner._finish_omni_aux_output(None, SimpleNamespace()) is None diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 0cd0a1ab3..847e16154 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -3,6 +3,7 @@ import contextlib from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -22,6 +23,23 @@ class _StopExecutionError(Exception): """Halts execute_model at the model boundary.""" +@pytest.mark.parametrize("randomize_inputs", [False, True]) +def test_generation_profile_accepts_upstream_randomize_inputs(randomize_inputs): + """Release workers can pass the profiling flag to codec-stage runners.""" + runner = object.__new__(GPUGenerationModelRunner) + runner.supports_mm_inputs = False + runner.max_num_tokens = 128 + runner.encoder_cache = {"temporary": object()} + runner._dummy_run = Mock(return_value=(torch.zeros(1), None)) + runner._sync_device = Mock() + + runner.profile_run(randomize_inputs=randomize_inputs) + + runner._dummy_run.assert_called_once_with(128, is_profile=True, randomize_inputs=randomize_inputs) + runner._sync_device.assert_called_once_with() + assert runner.encoder_cache == {} + + class _DummyInputBatch: def __init__(self, num_reqs: int = 1): self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)] diff --git a/tests/worker/test_omni_v2_profile_dispatch.py b/tests/worker/test_omni_v2_profile_dispatch.py index 4ff6fbfd4..14d86924a 100644 --- a/tests/worker/test_omni_v2_profile_dispatch.py +++ b/tests/worker/test_omni_v2_profile_dispatch.py @@ -15,8 +15,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size): runner = object.__new__(OmniGPUModelRunner) runner.cudagraph_manager = None - runner.dp_size = dp_size - runner.dp_rank = 1 + runner.parallel_config = SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=1) + runner.dp_size, runner.dp_rank = dp_size, 1 import vllm.v1.worker.gpu.dp_utils as dp_utils @@ -52,6 +52,7 @@ def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size): return expected runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch) + runner.parallel_config.data_parallel_size = 1 runner.dp_size = 1 batch_desc, _ = runner._dispatch_batch_descriptor( num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8 diff --git a/tests/worker_v2/test_capture_query_bound_gpu.py b/tests/worker_v2/test_capture_query_bound_gpu.py index c341a575b..2b4c37878 100644 --- a/tests/worker_v2/test_capture_query_bound_gpu.py +++ b/tests/worker_v2/test_capture_query_bound_gpu.py @@ -44,8 +44,8 @@ def _check_prefill(): spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16) kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) state = OmniModelState.__new__(OmniModelState) + state.model_config = SimpleNamespace(max_model_len=256) state.vllm_config = c - state.max_model_len = 256 state.supports_mm_inputs = False with set_current_vllm_config(c): group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0) diff --git a/tests/worker_v2/test_decode_query_bound_gpu.py b/tests/worker_v2/test_decode_query_bound_gpu.py index 114e9f343..9089fb76b 100644 --- a/tests/worker_v2/test_decode_query_bound_gpu.py +++ b/tests/worker_v2/test_decode_query_bound_gpu.py @@ -54,7 +54,7 @@ def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path): kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) state = OmniModelState.__new__(OmniModelState) state.vllm_config = c - state.max_model_len = 2048 + state.model_config = c.model_config state.supports_mm_inputs = False with set_current_vllm_config(c): group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0) diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index cd0c09877..fe89df3a2 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -11,10 +11,13 @@ from unittest.mock import MagicMock import numpy as np import pytest import torch -from vllm.v1.outputs import RoutedExpertsTensors +from vllm.distributed.aux_output_connector.connector import AuxRequestOutput +from vllm.sampling_params import SamplingParams from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors +from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner +from tests.helpers.mark import hardware_test from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot @@ -41,7 +44,7 @@ class _FakeEvent: def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput: rid2idx = {rid: i for i, rid in enumerate(req_ids)} mro = omni_ar_model_runner.OmniModelRunnerOutput(list(req_ids), rid2idx, None, prompt_logprobs_dict={}) - sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, None) + sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, torch.tensor([1]), torch.tensor([0])) kwargs = dict(model_runner_output=mro, sampler_output=sampler_output) kwargs.update(num_sampled_tokens=torch.tensor([1] * len(req_ids)), copy_event=_FakeEvent()) kwargs.update(main_stream=_FakeStream(), copy_stream=_FakeStream()) @@ -64,23 +67,38 @@ def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_widt counts=torch.tensor([2, 0]), vocab_size=4, ) - sampler_output = SamplerOutput(torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), None, masks) - routed = RoutedExpertsTensors(torch.tensor([[[2, 3]], [[4, 5]]]), torch.tensor([7, 9])) + sampler_output = SamplerOutput( + torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), torch.tensor([0, 0]), masks + ) + copied: list[tuple[np.ndarray, np.ndarray]] = [] + class _FakePendingAuxOutput: + def enqueue_cpu_copy(self, *, num_sampled, num_rejected) -> None: + copied.append((num_sampled, num_rejected)) + + def process_output(self) -> dict[str, AuxRequestOutput]: + return { + "decode": AuxRequestOutput(token_start=0, rows=np.array([[2, 3]], dtype=np.uint8)), + "prefill": AuxRequestOutput(token_start=0, rows=np.array([[4, 5]], dtype=np.uint8)), + } + + pending = _FakePendingAuxOutput() output = _async_output( req_ids=["decode", "prefill"], sampler_output=sampler_output, num_sampled_tokens=torch.tensor([1, 0]), copy_event=None, # None → constructor builds the default blocking event - routed_experts=routed, + pending_aux_output=pending, ).get_output() assert event_kwargs == [{"blocking": True}] # blocking event by default assert output.sampled_token_ids == [[2], []] assert output.sampling_masks.to_nested_list() == [[0, 2], []] np.testing.assert_array_equal(output.sampling_masks.offsets, [0, 2, 2]) - np.testing.assert_array_equal(output.routed_experts.routing_data, [[[2, 3]], [[4, 5]]]) - np.testing.assert_array_equal(output.routed_experts.slot_mapping, [7, 9]) + np.testing.assert_array_equal(output.aux_output_connector_output["decode"].rows, [[2, 3]]) + np.testing.assert_array_equal(output.aux_output_connector_output["prefill"].rows, [[4, 5]]) + assert output.aux_output_connector_output["decode"].token_start == 0 np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2]) + assert len(copied) == 1 and copied[0][0].tolist() == [1, 0] @pytest.mark.parametrize("needs_history", [False, True]) @@ -88,19 +106,26 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) runner = OmniARModelRunner.__new__(OmniARModelRunner) input_batch = SimpleNamespace(req_ids=["req"], num_reqs=1, seq_lens=torch.tensor([3])) input_batch.idx_mapping, input_batch.query_start_loc = torch.tensor([0]), torch.tensor([0, 1]) + input_batch.idx_mapping_np = np.array([0]) + input_batch.num_computed_prefill_tokens_np = np.array([0]) + input_batch.num_scheduled_tokens = np.array([3]) + input_batch.prefill_len_np = np.array([3]) + input_batch.query_start_loc_np = np.array([0, 3]) state = SimpleNamespace(input_batch=input_batch, hidden_states=torch.zeros(1, 2)) - state.finished_req_ids, state.ec_connector_output, state.routed_experts = {"finished"}, None, None + state.finished_req_ids, state.ec_connector_output = {"finished"}, None runner.execute_model_state = state runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False + runner.aux_output_connector = None runner.model_config = SimpleNamespace(async_chunk=False) runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text")) - runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(torch.zeros(1, 2), None))) + text_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(text_hidden, None))) runner.model_state.intermediate_buffer = SimpleNamespace(buffers={0: {"global_request_id": "global-req"}}) runner.req_states = SimpleNamespace(req_id_to_index={"req": 0}) runner.req_states.all_token_ids = SimpleNamespace(gpu=torch.tensor([[1]])) runner.req_states.num_computed_tokens = SimpleNamespace(gpu=torch.tensor([0])) - runner.req_states.prompt_len = SimpleNamespace(np=np.array([1]), gpu=torch.tensor([1])) + runner.req_states.prompt_len = SimpleNamespace(np=np.array([3]), gpu=torch.tensor([3])) runner.main_stream = runner.output_copy_stream = MagicMock() runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = MagicMock() sampler_out = (SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), MagicMock(), MagicMock()) @@ -114,13 +139,17 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) yield sampling_active = False - runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history) + runner.model = SimpleNamespace( + compute_logits=lambda hidden: hidden, logitsprocs_need_output_token_ids=needs_history + ) runner.model.mrv2_sampling_context = sampling_context runner.sampler = None runner.sample = MagicMock(return_value=sampler_out) runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail() logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {}) - runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=logprobs_mock) + runner.prompt_logprobs_worker = PromptLogprobsWorker(1, torch.device("cpu"), logprobs_mode="raw_logits") + runner.prompt_logprobs_worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[1, 0])) + runner.prompt_logprobs_worker.compute_prompt_logprobs = logprobs_mock runner.postprocess_sampled, connector_output = MagicMock(), object() def post_forward(finished_req_ids): @@ -135,10 +164,67 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) assert runner.sample_tokens(None) is mock_out built = omni_ar_model_runner.OmniAsyncOutput.call_args.kwargs["model_runner_output"] assert built.kv_connector_output is connector_output + torch.testing.assert_close(built.prompt_token_id_logprobs_dict["req"], torch.tensor([[2.0, 1.0], [4.0, 3.0]])) assert runner._resolve_global_request_id("req") == "global-req" # from the intermediate buffer assert runner._resolve_global_request_id("unknown") == "unknown" # fallback to the local id +def test_async_output_preserves_fixed_token_scores_on_host(monkeypatch) -> None: + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]]) + mro = omni_ar_model_runner.OmniModelRunnerOutput( + req_ids=["req-0"], + req_id_to_index={"req-0": 0}, + sampled_token_ids=None, + prompt_logprobs_dict={}, + prompt_token_id_logprobs_dict={"req-0": scores}, + ) + output = _async_output(model_runner_output=mro).get_output() + assert output.prompt_token_id_logprobs_dict["req-0"].device.type == "cpu" + torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req-0"], scores) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_chunked_fixed_token_scores_are_copied_from_cuda() -> None: + device = torch.device("cuda") + hidden = torch.tensor([[1.0, 2.0, 3.0], [3.0, 1.0, 2.0], [2.0, 3.0, 1.0], [4.0, 5.0, 6.0]], device=device) + worker = PromptLogprobsWorker(1, device) + worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[2, 0], prompt_logprob_start=1)) + batch = SimpleNamespace( + req_ids=["req"], + idx_mapping_np=np.array([0]), + num_computed_prefill_tokens_np=np.array([0]), + num_scheduled_tokens=np.array([2]), + prefill_len_np=np.array([4]), + query_start_loc_np=np.array([0, 2]), + ) + assert worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[:2], batch, np.array([4])) == {} + batch.num_computed_prefill_tokens_np = np.array([2]) + scores = worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[2:], batch, np.array([4])) + mro = omni_ar_model_runner.OmniModelRunnerOutput( + req_ids=["req"], + req_id_to_index={"req": 0}, + sampled_token_ids=None, + prompt_logprobs_dict={}, + prompt_token_id_logprobs_dict=scores, + ) + num_sampled = torch.tensor([1], device=device, dtype=torch.int32) + sampler_output = SamplerOutput( + torch.tensor([[1]], device=device), None, None, num_sampled, torch.zeros_like(num_sampled) + ) + output = OmniAsyncOutput( + model_runner_output=mro, + sampler_output=sampler_output, + num_sampled_tokens=num_sampled, + main_stream=torch.cuda.current_stream(), + copy_stream=torch.cuda.Stream(), + ).get_output() + + expected = hidden[1:3].log_softmax(dim=-1)[:, [2, 0]].cpu() + assert output.prompt_token_id_logprobs_dict["req"].device.type == "cpu" + torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req"], expected) + + def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: runner = OmniARModelRunner.__new__(OmniARModelRunner) runner.model_config = SimpleNamespace(async_chunk=True) @@ -284,7 +370,13 @@ def test_async_output_slices_request_payloads_with_graph_padding( batch.num_tokens_after_padding = padded_total output = _async_output( req_ids=[f"req-{i}" for i in range(len(lengths))], - sampler_output=SamplerOutput(torch.ones(len(lengths), 1, dtype=torch.long), None, None, None), + sampler_output=SamplerOutput( + torch.ones(len(lengths), 1, dtype=torch.long), + None, + None, + torch.ones(len(lengths), dtype=torch.long), + torch.zeros(len(lengths), dtype=torch.long), + ), text_hidden=hidden, multimodal_outputs={"codes": {"audio": codes, "ref": refs}}, input_batch=batch, @@ -585,7 +677,9 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, result = _async_output( req_ids=["audio", "partial"], - sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None), + sampler_output=SamplerOutput( + torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0]) + ), num_sampled_tokens=torch.tensor([1, 0]), multimodal_outputs={"snapshot": snapshot}, input_batch=batch, @@ -710,7 +804,7 @@ def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocke extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None result = _async_output( req_ids=["frame", "eos", "prefill"], - sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None), + sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None, None), num_sampled_tokens=torch.tensor([1, 1, 1]), multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}}, input_batch=batch, @@ -753,7 +847,9 @@ def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming): result = _async_output( req_ids=["audio", "partial"], - sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None), + sampler_output=SamplerOutput( + torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0]) + ), num_sampled_tokens=torch.tensor([1, 0]), multimodal_outputs={"snapshot": snapshot}, input_batch=batch, diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py index 2b12837e4..df1d9f3ca 100644 --- a/tests/worker_v2/test_omni_generation_model_runner.py +++ b/tests/worker_v2/test_omni_generation_model_runner.py @@ -4,12 +4,14 @@ """OmniGenerationModelRunner contracts: empty-step lifecycle, output partition, CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling.""" +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock import numpy as np import pytest import torch +from vllm.distributed.aux_output_connector.connector import AuxRequestOutput from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.outputs import OmniModelRunnerOutput @@ -18,6 +20,17 @@ from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModel pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("randomize_inputs", [False, True]) +def test_profile_run_accepts_release_flag_without_running_codec(monkeypatch, randomize_inputs): + runner = object.__new__(OmniGenerationModelRunner) + runner.model = MagicMock(side_effect=AssertionError("profiling ran codec model")) + synchronize = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + runner.profile_run(randomize_inputs=randomize_inputs) + synchronize.assert_called_once_with() + runner.model.assert_not_called() + + class _FakeStagedField: def __init__(self, data): self.np = data @@ -45,6 +58,7 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10): runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None) runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None)) runner.check_ep_fault = False + runner.aux_output_connector = None runner.req_states = MagicMock( prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)), num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)), @@ -52,7 +66,8 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10): return runner -def test_control_only_step_keeps_lifecycle_and_skips_input_construction(): +@pytest.mark.parametrize("aux_enabled", [False, True]) +def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enabled): runner = object.__new__(OmniGenerationModelRunner) order = [] for name in ( @@ -68,8 +83,13 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction(): runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output) runner._merge_ec_connector_no_forward = lambda _s, value: value runner._attach_native_data_plane_signals = lambda value: value + metadata = object() + runner.aux_output_connector = SimpleNamespace(begin_step=MagicMock()) if aux_enabled else None scheduler_output = SimpleNamespace( - total_num_scheduled_tokens=0, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace(req_ids=[]) + total_num_scheduled_tokens=0, + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(req_ids=[]), + aux_output_connector_metadata=metadata, ) assert runner.execute_model(scheduler_output) is output assert order == [ @@ -78,6 +98,9 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction(): "free_states", "_apply_block_table_staged_writes_if_available", ] + if aux_enabled: + assert runner.aux_output_connector is not None + runner.aux_output_connector.begin_step.assert_called_once_with(metadata) def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state(): @@ -184,6 +207,72 @@ def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypa assert output_req_ids == ["req-0"] +def test_generation_aux_output_copies_and_materializes_without_sampling(monkeypatch): + from vllm_omni.worker_v2 import omni_generation_model_runner as module + + events = [] + + class PendingOutput: + def enqueue_cpu_copy(self, *, num_sampled, num_rejected): + np.testing.assert_array_equal(num_sampled, [0]) + np.testing.assert_array_equal(num_rejected, [0]) + events.append("copy") + + def process_output(self): + assert events == ["copy", "record", "synchronize"] + return {"req-0": AuxRequestOutput(0, np.array([[2, 3]], dtype=np.uint8))} + + class CopyEvent: + def record(self, stream): + events.append("record") + + def synchronize(self): + events.append("synchronize") + + monkeypatch.setattr(torch.cuda, "Event", lambda **kw: CopyEvent()) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(module, "_async_copy_mm", lambda *a, **kw: {}) + output = OmniModelRunnerOutput(["req-0"], {"req-0": 0}, [[]]) + result = module.OmniGenerationAsyncOutput( + model_runner_output=output, + multimodal_outputs={}, + num_reqs=1, + main_stream=object(), + copy_stream=SimpleNamespace(wait_stream=lambda stream: None), + pending_aux_output=PendingOutput(), + ).get_output() + + assert result.sampled_token_ids == [[]] + np.testing.assert_array_equal(result.aux_output_connector_output["req-0"].rows, [[2, 3]]) + + +def test_pending_aux_output_forces_async_delivery_without_cuda_payload(monkeypatch): + from vllm_omni.worker_v2 import omni_generation_model_runner as module + + runner = _make_runner(OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={})) + runner.device = SimpleNamespace(type="cuda") + runner.main_stream = object() + runner.output_copy_stream = object() + pending = object() + runner.aux_output_connector = SimpleNamespace(prepare_output=MagicMock(return_value=pending)) + runner._reserve_native_data_plane_outputs = MagicMock() + runner._release_generation_slots = MagicMock() + runner._finalize_native_data_plane_output = MagicMock() + captured = {} + + class AsyncOutput: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(module, "OmniGenerationAsyncOutput", AsyncOutput) + input_batch = runner._gen_input_batch + result = runner.sample_tokens() + + assert isinstance(result, AsyncOutput) + runner.aux_output_connector.prepare_output.assert_called_once_with(input_batch) + assert captured["pending_aux_output"] is pending + + @pytest.mark.parametrize("exact", [True, False]) def test_exact_input_shape_models_reject_padded_batches(exact): # Qwen3-Omni Code2Wav splits input_ids by seq_token_counts (#6712). diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py index dd4c760c6..803d7581a 100644 --- a/tests/worker_v2/test_omni_gpu_model_runner.py +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -13,6 +13,7 @@ from vllm import SamplingParams from vllm.config.compilation import CUDAGraphMode from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.worker.gpu.model_runner import GPUModelRunner +from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm_omni.config.model import OmniModelConfig @@ -95,7 +96,8 @@ def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker): @pytest.mark.parametrize("output_form", ["tuple", "omni"]) -def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): +@pytest.mark.parametrize("profile_only", [False, True]) +def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form, profile_only): runner = object.__new__(OmniGPUModelRunner) hidden = torch.ones(1, 2) @@ -117,16 +119,21 @@ def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock()) runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc") - def assert_unwrapped(_self): + def assert_unwrapped(_self, *, profile_only: bool = False): assert torch.equal(runner.model.forward(), hidden) # forward unwrapped during capture + assert profile_only is expected_profile_only return 3 + expected_profile_only = profile_only with patch.object(GPUModelRunner, "capture_model", assert_unwrapped): - assert runner.capture_model() == 3 + assert runner.capture_model(profile_only=profile_only) == 3 assert runner.model.forward is original_forward # restored after capture assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]} - runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor) + if profile_only: + runner.model_state.capture_mtp_graphs.assert_not_called() + else: + runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor) @pytest.mark.parametrize( @@ -163,6 +170,28 @@ def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect assert state is upstream.return_value +@pytest.mark.parametrize("omni_stage", [False, True]) +@pytest.mark.parametrize("specialized", [False, True]) +def test_auxiliary_default_state_has_omni_lifecycle_but_preserves_upstream_states(monkeypatch, omni_stage, specialized): + class SpecializedState(DefaultModelState): + pass + + state_cls = SpecializedState if specialized else DefaultModelState + upstream_state = object.__new__(state_cls) + monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", lambda *_args: upstream_state) + monkeypatch.setattr(OmniModelState, "__init__", lambda *_args: None) + model_config = object.__new__(OmniModelConfig) if omni_stage else SimpleNamespace() + state = init_omni_model_state( + SimpleNamespace(model_config=model_config), SimpleNamespace(), None, torch.device("cpu") + ) + if omni_stage and not specialized: + assert isinstance(state, OmniModelState) + assert callable(state.run_preprocess) + assert callable(state.postprocess_model_output) + else: + assert state is upstream_state + + def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch): runner = _make_runner() calls = [] @@ -209,7 +238,7 @@ def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves(): runner.model_state = SimpleNamespace() captured = {} - def capture(_self): + def capture(_self, *, profile_only=False): # The graph manager stores (hidden, leaves) in its aux buffers. assert runner.use_aux_hidden_state_outputs is True hidden, leaves = runner.model.forward(4) @@ -262,7 +291,9 @@ def test_full_graph_output_without_contract_is_hidden_only(): @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"]) -def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): +@pytest.mark.parametrize("randomize_inputs", [False, True]) +@pytest.mark.parametrize("is_profile", [False, True]) +def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inputs, is_profile, monkeypatch): from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_runner import ExecuteModelState @@ -288,6 +319,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): runner.model_state.prepare_inputs.return_value = {} runner._omni_data_plane = object() runner.supports_mm_inputs = False + runner.vocab_size = 32 runner.lora_config = None runner.is_encoder_decoder = False runner.eplb = MagicMock() @@ -308,7 +340,15 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): cg_mode=CUDAGraphMode.NONE, num_reqs=1, num_tokens=1, num_active_loras=0, max_query_len=1 ) runner._dispatch_batch_descriptor = MagicMock(return_value=(batch_desc, None)) - monkeypatch.setattr(InputBatch, "make_dummy", lambda *args, **kwargs: input_batch) + make_dummy = MagicMock(return_value=input_batch) + monkeypatch.setattr(InputBatch, "make_dummy", make_dummy) + randomized = [] + + def randomize(tensor, low, high): + randomized.append((low, high)) + return tensor.fill_(7) + + monkeypatch.setattr(torch.Tensor, "random_", randomize) monkeypatch.setattr("vllm_omni.worker_v2.omni_model_runner.build_slot_mappings_by_layer", lambda *args: {}) for module in ("omni_model_runner", "omni_generation_model_runner"): monkeypatch.setattr(f"vllm_omni.worker_v2.{module}.set_forward_context", lambda *args, **kwargs: nullcontext()) @@ -318,7 +358,20 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): # Upstream _dummy_run always supplies valid_dummy_state_slots, and the # result must use the real upstream constructor rather than a mocked state. - assert runner.execute_model(scheduled, dummy_run=True, valid_dummy_state_slots=True) is None + assert ( + runner.execute_model( + scheduled, + dummy_run=True, + valid_dummy_state_slots=True, + randomize_inputs=randomize_inputs, + is_profile=is_profile, + ) + is None + ) + assert make_dummy.call_args.kwargs["is_padding"] is not is_profile + should_randomize = randomize_inputs and runner_kind != "generation" + assert randomized == ([(0, 32)] if should_randomize else []) + assert input_batch.input_ids.tolist() == ([7] if should_randomize else [1]) assert isinstance(runner.execute_model_state, ExecuteModelState) assert runner.execute_model_state.input_batch is input_batch assert runner.execute_model_state.cudagraph_stats is None diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 05adbf413..3e3a94fc2 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -22,6 +22,16 @@ from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _m pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_prepare_attn_forwards_release_model_specific_metadata(): + state = object.__new__(OmniModelState) + metadata = object() + output = object() + batch = SimpleNamespace() + with patch.object(DefaultModelState, "prepare_attn", return_value=output) as prepare: + assert state.prepare_attn(batch, None, (), None, [], None, model_specific_attn_metadata=metadata) is output + assert prepare.call_args.kwargs["model_specific_attn_metadata"] is metadata + + class _DummyInputBatch: is_prefilling_np: np.ndarray req_ids: list[str] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 877de1d78..f7c4c3a6b 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -25,7 +25,6 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.utils import ( free_kv_blocks_in_physical_order, - omni_routed_experts_for_request, ) from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.serialization import deserialize_additional_information @@ -99,7 +98,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): core scheduling logic. """ - max_num_running_reqs: int + max_num_active_reqs: int def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool: model_config = self.vllm_config.model_config @@ -380,32 +379,32 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # them. Upstream vllm raises RuntimeError on this status; omni allows # async abort (e.g. client disconnect during TTS streaming) to leave # requests in the waiting/running queues temporarily. - waiting = getattr(self, "waiting") self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="ar") self._drop_aborted_queued_requests() self._resync_streaming_input_counter() - original_waiting = None + original_wait_queues = None if self._should_defer_waiting_admission(): - original_waiting = waiting + original_wait_queues = self.waiting, self.kv_holding_waiting self.waiting = create_request_queue(self.policy) + self.kv_holding_waiting = create_request_queue(self.policy) - original_max_num_running_reqs = self.max_num_running_reqs + original_max_num_active_reqs = self.max_num_active_reqs async_chunk_transport = self._async_chunk_transport_enabled() reserved_running_slots = ( self._get_async_chunk_reserved_running_slots() if async_chunk_transport and self.use_v2_model_runner else 0 ) if reserved_running_slots: - self.max_num_running_reqs = max(0, original_max_num_running_reqs - reserved_running_slots) + self.max_num_active_reqs = max(0, original_max_num_active_reqs - reserved_running_slots) try: scheduler_output = super().schedule(throttle_prefills) finally: - self.max_num_running_reqs = original_max_num_running_reqs - if original_waiting is not None: - deferred_waiting = list(self.waiting) - if deferred_waiting: - original_waiting.prepend_requests(deferred_waiting) - self.waiting = original_waiting + self.max_num_active_reqs = original_max_num_active_reqs + if original_wait_queues is not None: + original_waiting, original_kv_holding = original_wait_queues + original_waiting.prepend_requests(self.waiting) + original_kv_holding.prepend_requests(self.kv_holding_waiting) + self.waiting, self.kv_holding_waiting = original_wait_queues self._restore_omni_wait_queues() self._postprocess_omni_schedule_output( @@ -428,6 +427,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): sampled_token_ids = model_runner_output.sampled_token_ids logprobs = model_runner_output.logprobs prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + prompt_token_id_logprobs_dict = getattr(model_runner_output, "prompt_token_id_logprobs_dict", {}) num_scheduled_tokens = scheduler_output.num_scheduled_tokens pooler_outputs = model_runner_output.pooler_output mm_outputs = getattr(model_runner_output, "multimodal_outputs", None) @@ -704,8 +704,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if stopped: if self.chunk_transfer_adapter is not None: confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request) - if model_runner_output.routed_experts is not None: - routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) + if ( + self.aux_output_connector is not None + and (aux_output := model_runner_output.aux_output_connector_output) is not None + ): + routed_experts = self.aux_output_connector.take_output(request, aux_output) # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. @@ -778,6 +781,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id) has_stage_output = ( bool(new_token_ids) or mm_output is not None @@ -798,6 +802,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, + prompt_token_id_logprobs=prompt_token_id_logprobs, pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, @@ -813,6 +818,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors + assert prompt_token_id_logprobs is None if omits_chunk_transfer and inter_stage_output is not None: logger.warning( @@ -976,8 +982,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: self.num_waiting_for_streaming_input -= 1 session.status = RequestStatus.WAITING - if session in self.skipped_waiting: - self.skipped_waiting.remove_requests((session,)) + self.deferred_waiting.discard(session) + if session in self.kv_holding_waiting: + self.kv_holding_waiting.remove_requests((session,)) self._enqueue_waiting_request(session) if self.log_stats: @@ -1401,7 +1408,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): ``_free_blocks``) and skips the caller's input-coordinator cleanup. """ self.waiting.remove_requests((request,)) - self.skipped_waiting.remove_requests((request,)) + self.kv_holding_waiting.remove_requests((request,)) + self.deferred_waiting.discard(request) if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: self.num_waiting_for_streaming_input -= 1 request.status = RequestStatus.FINISHED_ERROR diff --git a/vllm_omni/core/sched/omni_cfg_ar_scheduler.py b/vllm_omni/core/sched/omni_cfg_ar_scheduler.py index f1a72ac81..7d562502e 100644 --- a/vllm_omni/core/sched/omni_cfg_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_cfg_ar_scheduler.py @@ -32,7 +32,7 @@ class OmniCFGARScheduler(OmniARScheduler): raise ValueError("Atomic CFG reservations require one full-attention KV group") self._capacity = self.kv_cache_manager.block_pool.num_gpu_blocks - 1 - self.kv_cache_manager.watermark_blocks pair_pages = 2 * -(-self.max_model_len // self.block_size) - if self._capacity < pair_pages or self.max_num_running_reqs < 2: + if self._capacity < pair_pages or min(self.max_num_running_reqs, self.max_num_active_reqs) < 2: raise ValueError("Atomic CFG requires KV capacity and max_num_seqs for two maximum-length branches") if self.max_num_scheduled_tokens < 2 * self.max_model_len: raise ValueError("Atomic CFG requires max_num_batched_tokens >= 2 * max_model_len") @@ -81,7 +81,7 @@ class OmniCFGARScheduler(OmniARScheduler): admitted = create_request_queue(self.policy) pending = {request.request_id: request for request in original_waiting} pages = self._capacity - sum(self._reserved_pages(request) for request in self.running) - slots = self.max_num_running_reqs - len(self.running) + slots = self.max_num_active_reqs - len(self.running) tokens = self.max_num_scheduled_tokens - len(self.running) chosen: set[str] = set() for request in original_waiting: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index b103776fa..b0c47b2fe 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -26,7 +26,6 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData -from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.outputs import OmniModelRunnerOutput @@ -275,6 +274,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.connector is not None: meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta + if self.aux_output_connector is not None: + scheduler_output.aux_output_connector_metadata = self.aux_output_connector.build_connector_meta( + scheduler_output, self.requests + ) # EC Connector: package metadata if self.ec_connector is not None: ec_meta = self.ec_connector.build_connector_meta(scheduler_output) @@ -330,7 +333,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ): # Downstream async-chunk stages receive the next segment from the # connector, not from an API StreamingUpdate. Enqueue them as - # schedulable before the base class can park them in skipped_waiting. + # schedulable before the base class can defer a streaming-input wait. request.status = RequestStatus.WAITING self._enqueue_waiting_request(request) return False @@ -342,7 +345,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): A generation batch slot belongs to one chunk, not the entire stream. Keep requests with outstanding output in ``running`` so their state cannot be reused before completion. Other live requests join the tail - of ``waiting`` without freeing their decoder state or KV blocks. + of the appropriate waiting queue without freeing decoder state or KV blocks. Preserve WAITING_FOR_CHUNK until the coordinator observes fresh input. Stateful codecs retain lifetime admission until their decoder state is released, even when no chunk is currently executing. @@ -360,7 +363,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): continue if request.status == RequestStatus.RUNNING: request.status = RequestStatus.WAITING - self.waiting.add_request(request) + self._enqueue_waiting_request(request) self.running = in_flight def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: @@ -391,6 +394,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Temporary queue: preserve waiting order while requests await input. skipped_waiting_requests = create_request_queue(self.policy) + skipped_kv_holding_requests = create_request_queue(self.policy) req_index = 0 self._drop_aborted_queued_requests() self._requeue_completed_native_chunks() @@ -502,7 +506,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Fast path selection and scheduling (treat all as diffusion requests, # independent of pooling_params) while ( - self.waiting + (self.kv_holding_waiting or self.waiting) and token_budget > 0 and len(num_scheduled_tokens) < execution_batch_size and ( @@ -510,38 +514,47 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if native_chunks and not self._retains_state_across_chunks else len(self.running) + reserved_running_slots ) - < self.max_num_running_reqs + < self.max_num_active_reqs and self._pause_state == PauseState.UNPAUSED ): - request = self.waiting.peek_request() + request_queue = self.kv_holding_waiting or self.waiting + skipped_requests = ( + skipped_kv_holding_requests if request_queue is self.kv_holding_waiting else skipped_waiting_requests + ) + request = request_queue.peek_request() if express and request.request_id in self._chunk_started: - self.waiting.pop_request() - skipped_waiting_requests.add_request(request) + request_queue.pop_request() + self.deferred_waiting.add(request) + skipped_requests.add_request(request) continue if native_chunks and request.num_in_flight_tokens > 0: # A restored waiting entry can still own an outstanding batch. # Do not let it block other ready streams or execute it twice. - self.waiting.pop_request() - skipped_waiting_requests.add_request(request) + request_queue.pop_request() + self.deferred_waiting.add(request) + skipped_requests.add_request(request) continue # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED ): # Pop the finished request from waiting queue and don't schedule it - self.waiting.pop_request() + request_queue.pop_request() + self.deferred_waiting.discard(request) continue # async_chunk: wait for the first upstream chunk (don't start with placeholders). if async_chunk_transport and len(request.prompt_token_ids) == 0: if self._is_done_receiving_chunks(request.request_id): if not _has_async_chunk_payload_to_run(request): - self.waiting.pop_request() + request_queue.pop_request() + self.deferred_waiting.discard(request) self._pending_finish_reqs.append(request) continue else: - self.waiting.pop_request() - skipped_waiting_requests.prepend_request(request) + request_queue.pop_request() + self.deferred_waiting.add(request) + skipped_requests.prepend_request(request) continue # Allocate all input tokens for the request in one shot @@ -557,11 +570,13 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ): required_tokens = 1 else: - self.waiting.pop_request() + request_queue.pop_request() if self._is_done_receiving_chunks(request.request_id): + self.deferred_waiting.discard(request) self._pending_finish_reqs.append(request) else: - skipped_waiting_requests.add_request(request) + self.deferred_waiting.add(request) + skipped_requests.add_request(request) continue num_new_tokens = min(required_tokens, token_budget) new_blocks = self.kv_cache_manager.allocate_slots( @@ -577,7 +592,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): break # Officially schedule this request - request = self.waiting.pop_request() + request = request_queue.pop_request() + self.deferred_waiting.discard(request) self.running.append(request) if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) @@ -592,7 +608,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if first_chunk_express: self._chunk_started.update(num_scheduled_tokens) - # Return skipped waiting requests + # Return skipped requests to their original queue. + if skipped_kv_holding_requests: + self.kv_holding_waiting.prepend_requests(skipped_kv_holding_requests) if skipped_waiting_requests: self.waiting.prepend_requests(skipped_waiting_requests) @@ -936,8 +954,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): 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) + if ( + self.aux_output_connector is not None + and (aux_output := model_runner_output.aux_output_connector_output) is not None + ): + routed_experts = self.aux_output_connector.take_output(request, aux_output) finish_reason = request.get_finished_reason() finished = self._handle_stopped_request(request) is_segment_finished = not finished diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 105b79a00..f8e8bcb40 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -111,6 +111,56 @@ else: _SchedulerMixinBase = object +class _OmniWaitingQueue: + """Let one input-gating pass visit both upstream waiting queues. + + Requests holding KV blocks retain upstream admission priority when input + gates park, restore or requeue them. Connector state is consumed once per + scheduler tick, including when both physical queues contain requests. + """ + + def __init__(self, scheduler: Any) -> None: + self.scheduler = scheduler + + def __iter__(self) -> Iterator[Request]: + for request_queue in (self.scheduler.kv_holding_waiting, self.scheduler.waiting): + for request in request_queue: + # Upstream owns grammar, remote-KV and streaming-update waits. + # Omni input readiness must not overwrite those statuses. + if not self.scheduler._is_blocked_waiting_status(request.status): + yield request + + def add_request(self, request: Request) -> None: + self.scheduler._enqueue_waiting_request(request) + + def prepend_requests(self, requests: Iterable[Request]) -> None: + holding: list[Request] = [] + fresh: list[Request] = [] + for request in requests: + if self.scheduler._is_blocked_waiting_status(request.status): + self.scheduler.deferred_waiting.add(request) + (holding if self.scheduler._holds_kv_blocks(request) else fresh).append(request) + self.scheduler.kv_holding_waiting.prepend_requests(holding) + self.scheduler.waiting.prepend_requests(fresh) + + def remove_request(self, request: Request) -> None: + # Input metadata can change the computed frontier before removal, so + # choose the queue by membership rather than the request's new state. + if request in self.scheduler.kv_holding_waiting: + self.scheduler.kv_holding_waiting.remove_request(request) + else: + self.scheduler.waiting.remove_request(request) + self.scheduler.deferred_waiting.discard(request) + + remove = remove_request + + def remove_requests(self, requests: Iterable[Request]) -> None: + requests = set(requests) + self.scheduler.kv_holding_waiting.remove_requests(requests) + self.scheduler.waiting.remove_requests(requests) + self.scheduler.deferred_waiting.difference_update(requests) + + class OmniSchedulerMixin(_SchedulerMixinBase): """Shared scheduler helpers for omni-specific request handling.""" @@ -118,6 +168,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): vllm_config: VllmConfig requests: dict[str, Request] waiting: RequestQueue + kv_holding_waiting: RequestQueue + deferred_waiting: set[Request] running: list[Request] _pending_input_timeout_outputs: dict[str, tuple[int, OmniEngineCoreOutput]] @@ -243,8 +295,9 @@ class OmniSchedulerMixin(_SchedulerMixinBase): if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: self.num_waiting_for_streaming_input -= 1 session.status = RequestStatus.WAITING - if session in self.skipped_waiting: - self.skipped_waiting.remove_requests((session,)) + self.deferred_waiting.discard(session) + if session in self.kv_holding_waiting: + self.kv_holding_waiting.remove_requests((session,)) self._enqueue_waiting_request(session) if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) @@ -377,14 +430,14 @@ class OmniSchedulerMixin(_SchedulerMixinBase): ) if input_coordinator._async_chunk: input_coordinator.process_pending_chunks( - self.waiting, + _OmniWaitingQueue(self), self.running, chunk_ready_req_ids, chunk_finished_req_ids, ) else: input_coordinator.process_pending_full_payload_inputs( - self.waiting, + _OmniWaitingQueue(self), stage_recv_req_ids, ) @@ -394,7 +447,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): self._process_pending_input_timeouts() if self.chunk_transfer_adapter: self.chunk_transfer_adapter.process_pending_chunks( - self.waiting, + _OmniWaitingQueue(self), self.running, scheduler_requests=self.requests, ) @@ -426,12 +479,12 @@ class OmniSchedulerMixin(_SchedulerMixinBase): """Restore requests temporarily parked by Omni input gates.""" if self.chunk_transfer_adapter: self.chunk_transfer_adapter.restore_queues( - self.waiting, + _OmniWaitingQueue(self), self.running, scheduler_requests=self.requests, ) if self.input_coordinator: - self.input_coordinator.restore_queues(self.waiting, self.running) + self.input_coordinator.restore_queues(_OmniWaitingQueue(self), self.running) def _process_pending_input_timeouts(self) -> None: """Force-fail requests waiting on the full-payload coordinator too long. @@ -669,6 +722,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): finish_reason: Any = None, new_logprobs: Any = None, new_prompt_logprobs_tensors: Any = None, + prompt_token_id_logprobs: Any = None, pooling_output: Any = None, multimodal_output: Any = None, stop_reason: Any = None, @@ -693,6 +747,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=new_prompt_logprobs_tensors, + prompt_token_id_logprobs=prompt_token_id_logprobs, pooling_output=pooling_output, pooling_output_payload=pooling_output_payload, multimodal_output=multimodal_output, @@ -834,7 +889,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): ) } self.waiting.remove_requests(removable) - self.skipped_waiting.remove_requests(removable) + self.kv_holding_waiting.remove_requests(removable) + self.deferred_waiting.difference_update(removable) def _resume_downstream_chunk_receiver(self, request: Request) -> None: """Resume duplex connector polling without an external update.""" @@ -850,7 +906,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): return self.num_waiting_for_streaming_input -= 1 request.status = RequestStatus.WAITING - self.skipped_waiting.remove_requests((request,)) + self.deferred_waiting.discard(request) + self.kv_holding_waiting.remove_requests((request,)) self._enqueue_waiting_request(request) def _aggregate_kv_connector_stats( @@ -973,10 +1030,10 @@ class OmniSchedulerMixin(_SchedulerMixinBase): request_ids = tuple(request_ids) target_request_ids = set(request_ids) - skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())} + deferred_waiting_ids = {r.request_id for r in self.deferred_waiting} pre_adapter_streaming_wait_ids = { rid - for rid in target_request_ids & skipped_waiting_ids + for rid in target_request_ids & deferred_waiting_ids if (req := self.requests.get(rid)) is not None and ( req.status == RequestStatus.WAITING_FOR_STREAMING_REQ @@ -1038,7 +1095,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): status is not ``RUNNING``, set it to ``RUNNING``; symmetrically flip ``RUNNING → WAITING`` when the request is actually in ``self.waiting``. A resumable segment stop still held in - ``self.skipped_waiting`` is restored to + ``self.kv_holding_waiting`` is restored to ``WAITING_FOR_STREAMING_REQ`` so upstream also balances its paused-session counter. Because adapter cleanup may restore a connector-owned request's prior status first, ``finish_requests`` @@ -1075,8 +1132,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): return running_ids = {r.request_id for r in self.running} - waiting_ids = {r.request_id for r in self.waiting} - skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())} + waiting_ids = {r.request_id for queue in (self.waiting, self.kv_holding_waiting) for r in queue} + deferred_waiting_ids = {r.request_id for r in self.deferred_waiting} for rid in ids_to_align: req = self.requests.get(rid) @@ -1098,7 +1155,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): ) if req.is_finished() and not (resumable_segment_stop or streaming_wait): continue - if rid in skipped_waiting_ids and streaming_wait: + if rid in deferred_waiting_ids and streaming_wait: req.status = RequestStatus.WAITING_FOR_STREAMING_REQ elif rid in running_ids and req.status != RequestStatus.RUNNING: req.status = RequestStatus.RUNNING @@ -1161,10 +1218,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase): self.running[:] = [req for req in self.running if keep_running(req)] def _drop_aborted_queued_requests(self) -> None: - for request_queue in (self.waiting, self.skipped_waiting): + for request_queue in (self.waiting, self.kv_holding_waiting): aborted = [req for req in request_queue if req.status == RequestStatus.FINISHED_ABORTED] if aborted: request_queue.remove_requests(aborted) + self.deferred_waiting.difference_update(aborted) self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED] def _resync_streaming_input_counter(self) -> None: @@ -1173,7 +1231,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): return parked = sum( 1 - for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None)) + for queue in (getattr(self, "waiting", None), getattr(self, "kv_holding_waiting", None)) if queue for request in queue if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ diff --git a/vllm_omni/core/sched/utils.py b/vllm_omni/core/sched/utils.py index 12cadae73..d346caebb 100644 --- a/vllm_omni/core/sched/utils.py +++ b/vllm_omni/core/sched/utils.py @@ -2,29 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Shared utilities for omni schedulers.""" -import numpy as np from vllm.v1.core.kv_cache_manager import KVCacheManager -from vllm.v1.outputs import RoutedExpertsLists from vllm.v1.request import Request -def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None: - """Extract per-request routed experts from RoutedExpertsLists using slot_mapping. - - Matches upstream RoutedExpertsManager.get() pattern — filters routing_data - rows whose slot_mapping entries belong to this request's block_table. - """ - if routed_experts is None: - return None - slots = getattr(request, "block_table", None) - if slots is None: - return None - slot_set = set(slots) - mask = np.isin(routed_experts.slot_mapping, list(slot_set)) - data = routed_experts.routing_data[mask] - return data if data.size > 0 else None - - def free_kv_blocks_in_physical_order(manager: KVCacheManager, request: Request) -> None: """Return non-cached request blocks in increasing physical order. diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index 09c85f9ab..b8bc67af3 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -64,6 +64,7 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, spli from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.component_config import resolve_component_quant_config +from vllm_omni.quantization.fp8_config import DiffusionFp8Config logger = init_logger(__name__) @@ -322,7 +323,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports """Load MLLM through HF, or prepare it for deferred online FP8 loading.""" config = None if quant_config is not None: - if not isinstance(quant_config, Fp8Config): + if not isinstance(quant_config, (Fp8Config, DiffusionFp8Config)): raise ValueError( "Boogu MLLM only supports FP8 quantization. Set mllm to null to disable online quantization." ) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index c36d6d986..d7cb082e4 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -123,6 +123,8 @@ class OmniEngineCoreRequest(EngineCoreRequest): trace_headers=request.trace_headers, resumable=request.resumable, external_req_id=request.external_req_id, + session_id=request.session_id, + kv_hints=request.kv_hints, reasoning_ended=request.reasoning_ended, reasoning_parser_kwargs=request.reasoning_parser_kwargs, abort_immediately=request.abort_immediately, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index bad555704..a46fc49e2 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -270,6 +270,10 @@ class OmniEngineArgs(EngineArgs): ) validate_worker_omni_connector(self.worker_cls, needs_connector) super().__post_init__() + # The NPU runner implements the auxiliary connector on its legacy + # execution path; CUDA/ROCm Omni stages use the V2 implementation. + if self.aux_output_config.enabled and not self.use_v2_model_runner and not current_omni_platform.is_npu(): + raise ValueError("Auxiliary outputs require use_v2_model_runner=True for this Omni stage.") def _ensure_omni_models_registered(self): if hasattr(self, "_omni_models_registered"): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index c8ab8fcec..c86b882c8 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -14,6 +14,7 @@ from typing import Any, Literal, cast from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest +from vllm.v1.kv_hints import KvHintsEnvelope from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS from vllm_omni.engine import OmniEngineCoreRequest @@ -259,6 +260,7 @@ class AsyncOmniEngine(OmniEngineBase): data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, + kv_hints: KvHintsEnvelope | None = None, resumable: bool = False, message_type: Literal["add_request", "streaming_update"] = "add_request", ) -> StageSubmissionMessage: @@ -289,6 +291,8 @@ class AsyncOmniEngine(OmniEngineBase): request_artifact_dirs: list[str] = [] stage_type = self.stage_metadata[0].stage_type + if stage_type == "diffusion" and kv_hints is not None: + raise ValueError("kv_hints require an AR/LLM stage with a KV cache") output_prompt_text: Any = None _preprocess_ms = 0.0 if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): @@ -343,6 +347,7 @@ class AsyncOmniEngine(OmniEngineBase): priority=priority, data_parallel_rank=data_parallel_rank, resumable=resumable, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) except Exception: if preselected_stage0_replica is not None and self.stage_pools: @@ -518,6 +523,7 @@ class AsyncOmniEngine(OmniEngineBase): data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, + kv_hints: KvHintsEnvelope | None = None, resumable: bool = False, ) -> None: """Process stage-0 input locally, then send to the Orchestrator. @@ -543,6 +549,7 @@ class AsyncOmniEngine(OmniEngineBase): data_parallel_rank=data_parallel_rank, reasoning_ended=reasoning_ended, resumable=resumable, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) except BaseException: if isinstance(prompt, dict): @@ -599,6 +606,7 @@ class AsyncOmniEngine(OmniEngineBase): data_parallel_rank: int | None = None, reasoning_ended: bool | None = None, *, + kv_hints: KvHintsEnvelope | None = None, resumable: bool = False, ) -> None: """Async add_request API.""" @@ -617,6 +625,7 @@ class AsyncOmniEngine(OmniEngineBase): data_parallel_rank=data_parallel_rank, reasoning_ended=reasoning_ended, resumable=resumable, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) def add_streaming_update( diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index fad9a8f52..7f286fb66 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -751,7 +751,7 @@ class OrchestratorBase: stage_ids: list[int] = [] for pool in target_pools: clear_sender = pool.stage_type != "diffusion" and ( - method in ("reset_mm_cache", "sleep") + method in ("reset_mm_cache", "sleep", "release_kv_cache_memory") or (method == "pause_scheduler" and kwargs.get("clear_cache", args[1] if len(args) > 1 else True)) ) for replica_id in pool.live_replica_ids(): diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index a87654cf2..b4174f057 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -34,6 +34,7 @@ from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata if TYPE_CHECKING: + from vllm.renderers import BaseRenderer from vllm.v1.engine import EngineCoreOutput from vllm_omni.inputs.data import OmniTokensPrompt @@ -90,6 +91,7 @@ class StageEngineCoreClientBase(StageClientBase): coordinator: Any = None, client_count: int = 1, client_index: int = 0, + renderer: BaseRenderer | None = None, ) -> StageEngineCoreClient | DPLBStageEngineCoreClient: """Create the appropriate stage async client for the DP mode.""" parallel_config = vllm_config.parallel_config @@ -103,6 +105,7 @@ class StageEngineCoreClientBase(StageClientBase): coordinator=coordinator, client_count=client_count, client_index=client_index, + renderer=renderer, ) if parallel_config.data_parallel_size > 1 and not parallel_config.data_parallel_external_lb: @@ -118,6 +121,7 @@ class StageEngineCoreClientBase(StageClientBase): client_addresses: dict[str, str] | None = None, client_count: int = 1, client_index: int = 0, + renderer: BaseRenderer | None = None, *, metadata: StageMetadata | None = None, engine_manager: Any = None, @@ -183,13 +187,16 @@ class StageEngineCoreClientBase(StageClientBase): ) try: - super().__init__( + # Concrete clients place AsyncMPClient next in the cooperative MRO; + # the mixin's declared base exposes only object.__init__ to mypy. + super().__init__( # type: ignore[call-arg] vllm_config, executor_class, log_stats=log_stats, client_addresses=client_addresses, client_count=client_count, client_index=client_index, + renderer=renderer, ) if engine_manager is not None: self.resources.engine_manager = engine_manager diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index c84d9ed18..3f7ecbbe0 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -16,7 +16,8 @@ import signal from typing import Any import vllm.v1.engine.core as _vllm_engine_core_module -from vllm.logger import init_logger +from vllm.config import VllmConfig +from vllm.logger import configure_logging, init_logger from vllm.transformers_utils.config import ( maybe_register_config_serialize_by_value, ) @@ -195,6 +196,10 @@ class StageEngineCoreProc(EngineCoreProc): logging / metrics only. """ signal_callback: SignalCallback | None = None + vllm_config: VllmConfig = kwargs["vllm_config"] + if logging_config := getattr(vllm_config, "logging_config", None): + configure_logging(logging_config) + maybe_register_config_serialize_by_value() # Register vllm-omni reasoning parsers (e.g. step_audio) in this diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 176eb9dc0..b5d5bec6d 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -109,6 +109,7 @@ class StagePool: "pause_scheduler", "resume_scheduler", "sleep", + "release_kv_cache_memory", "wake_up", "reset_prefix_cache", "reset_encoder_cache", @@ -1361,7 +1362,7 @@ class StagePool: if timeout is not None: return await asyncio.wait_for(result, timeout=timeout) return await result - if method in self._CACHE_RESET_METHODS: + if method in self._CACHE_RESET_METHODS or method == "release_kv_cache_memory": return {"supported": False, "error": f"EngineCore helper {method}_async is unavailable"} return await client.collective_rpc_async( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 2a3253599..01c8ee588 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: from vllm.tokenizers import TokenizerLike from vllm.v1.engine import PauseMode from vllm.v1.engine.input_processor import InputProcessor + from vllm.v1.kv_hints import KvHintsEnvelope from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType @@ -125,6 +126,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): priority: int = 0, data_parallel_rank: int | None = None, session_id: str | None = None, + kv_hints: KvHintsEnvelope | None = None, reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, arrival_time: float | None = None, @@ -278,6 +280,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): arrival_time=wall_start_ts, lora_request=lora_request, first_chunk_submitted=first_chunk_submitted, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) await first_chunk_submitted else: @@ -289,6 +292,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): final_output_stage_ids=final_output_stage_ids, arrival_time=wall_start_ts, lora_request=lora_request, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) submit_ts = time.time() stage_first_ts = cast(list[float | None], req_state.metrics.stage_first_ts) @@ -365,6 +369,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): arrival_time: float, lora_request: Any = None, first_chunk_submitted: asyncio.Future[None] | None = None, + kv_hints: KvHintsEnvelope | None = None, ) -> asyncio.Task: """Submit a streaming input generator as incremental stage-0 updates.""" if not sampling_params_list: @@ -426,6 +431,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): arrival_time=arrival_time, lora_request=lora_request, resumable=True, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) ) has_submitted_first_chunk = True @@ -483,6 +489,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): arrival_time=arrival_time, lora_request=lora_request, resumable=False, + **({"kv_hints": kv_hints} if kv_hints is not None else {}), ) ) has_submitted_first_chunk = True @@ -907,6 +914,28 @@ class AsyncOmni(AsyncOmniBase, EngineClient): ) return all(self._coerce_stage_bool(result) for result in results) + async def release_kv_cache_memory( + self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S + ) -> None: + """Discard AR-stage KV memory after a completed generation pause. + + EngineCore checks that the scheduler is paused and all executor + memory is resident. Diffusion stages have no KV cache to release. + Restore KV memory with ``wake_up(tags=["kv_cache"])`` before resuming. + """ + ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids) + if not ar_stage_ids: + return + async with self._pause_cond: + if not self._paused: + raise RuntimeError("release_kv_cache_memory() requires a completed pause first") + await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0) + if 0 in ar_stage_ids: + await asyncio.wait_for(self._clear_frontend_mm_cache(), timeout=timeout) + await self._engine_core_rpc("release_kv_cache_memory", stage_ids=ar_stage_ids, timeout=timeout) + self._record_stage_sleep(ar_stage_ids, [CuMemTag.KV_CACHE.value]) + self._hold_admission_until_resume = True + async def sleep( self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort" ) -> list[OmniACK]: diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py index bb02fe961..33b870a26 100644 --- a/vllm_omni/inputs/mm_processor.py +++ b/vllm_omni/inputs/mm_processor.py @@ -7,12 +7,37 @@ from collections.abc import Mapping from typing import TypeVar from transformers import BatchFeature +from vllm.config.multimodal import BaseDummyOptions, MultiModalDummyOptions from vllm.multimodal.parse import MultiModalDataItems -from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + cached_encode, +) _I = TypeVar("_I", bound=BaseProcessingInfo) +class OmniDummyInputsBuilder(BaseDummyInputsBuilder[_I]): + """Restore Omni's processor-input hook on the vLLM 0.31 API.""" + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> ProcessorInputs: + dummy_text = self.get_dummy_text(mm_counts) + dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, MultiModalDummyOptions(mm_options or {})) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) + + tokenizer = self.info.ctx.tokenizer + dummy_prompt = [] if tokenizer is None else cached_encode(tokenizer, dummy_text, truncation=False) + return ProcessorInputs(prompt=dummy_prompt, mm_data_items=dummy_mm_items) + + class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): """Share data/passthrough handling, leaving model processing in its owner. @@ -23,17 +48,37 @@ class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> ProcessorInputs: + if isinstance(self.dummy_inputs, OmniDummyInputsBuilder): + return self.dummy_inputs.get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + return super().get_dummy_inputs(seq_len, mm_counts, mm_options) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: - kwargs = dict(hf_processor_mm_kwargs) + kwargs = dict(hf_kwargs) prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None) if prompt is None: prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough = self._get_hf_mm_data(valid_items) + # Raw processor/passthrough data (upstream removed `_get_hf_mm_data` in + # favour of `_get_hf_mm_inputs`, which also renames `audios` -> `audio` + # and injects a dummy text key). Omni's custom `_call_hf_processor` + # subclasses still expect the un-renamed keys, so extract them here. + mm_data: dict[str, object] = {} + passthrough: dict[str, object] = {} + for items in valid_items.values(): + if not items: + continue + mm_data.update(items.get_processor_data()) + passthrough.update(items.get_passthrough_data()) result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {}) result.update(passthrough) return result diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index dc4105b4d..e7337d92e 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -52,7 +52,7 @@ class OmniRenderer: def _routes_no_media_kwargs(prompt) -> bool: return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data") - def _process_singleton(self, prompt, *, skip_mm_cache: bool = False): + def _process_singleton(self, prompt, *, skip_mm_cache: bool = False, tok_params=None): renderer = cast(BaseRenderer, self) if self._routes_no_media_kwargs(prompt): inputs = renderer._process_multimodal( @@ -62,11 +62,14 @@ class OmniRenderer: mm_uuids=prompt.get("multi_modal_uuids"), skip_mm_cache=skip_mm_cache, ) + renderer._truncate_expanded_prompt(inputs, tok_params) else: - inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + inputs = cast(BaseRenderer, super())._process_singleton( + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) return self._with_omni_extras(inputs, prompt) - async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False): + async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False, tok_params=None): renderer = cast(BaseRenderer, self) if self._routes_no_media_kwargs(prompt): inputs = await renderer._process_multimodal_async( @@ -76,8 +79,11 @@ class OmniRenderer: mm_uuids=prompt.get("multi_modal_uuids"), skip_mm_cache=skip_mm_cache, ) + renderer._truncate_expanded_prompt(inputs, tok_params) else: - inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) + inputs = await cast(BaseRenderer, super())._process_singleton_async( + prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params + ) return self._with_omni_extras(inputs, prompt) diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 30dced9d0..320b3a0d1 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -259,19 +259,25 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + mm_data: dict[str, object] = {} + passthrough_data: dict[str, object] = {} + for items in valid_mm_items.values(): + if not items: + continue + mm_data.update(items.get_processor_data()) + passthrough_data.update(items.get_passthrough_data()) prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + processor = self.info.get_hf_processor(**hf_kwargs) has_image = "images" in mm_data has_img2img = "pixel_values_img2img" in mm_data processed_data = BatchFeature() if has_image: - image_kwargs = {**hf_processor_mm_kwargs, "is_img2img": False} + image_kwargs = {**hf_kwargs, "is_img2img": False} image_outputs = self.info.ctx.call_hf_processor( processor, {"text": prompt_text, "images": mm_data["images"]}, @@ -280,7 +286,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn processed_data.update(image_outputs) if has_img2img: - img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs) + img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_kwargs) img2img_kwargs["is_img2img"] = True img2img_outputs = self.info.ctx.call_hf_processor( processor, diff --git a/vllm_omni/model_executor/models/common/cfg_pairing.py b/vllm_omni/model_executor/models/common/cfg_pairing.py index ae4b3ac38..d104d6283 100644 --- a/vllm_omni/model_executor/models/common/cfg_pairing.py +++ b/vllm_omni/model_executor/models/common/cfg_pairing.py @@ -89,14 +89,12 @@ def _cfg_extra(request: Any, key: str) -> Any: def _wait_queues(scheduler: Any) -> list[Any]: - # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks - # structurally not-ready requests in a second ``skipped_waiting`` queue; - # pair-hold must cover both queues or a lone member parked there would - # be admitted without its partner. + # Pair holding must cover requests that retain KV blocks as well as + # fresh requests, or a lone member can be admitted without its partner. queues = [scheduler.waiting] - skipped = getattr(scheduler, "skipped_waiting", None) - if skipped is not None: - queues.append(skipped) + kv_holding = getattr(scheduler, "kv_holding_waiting", None) + if kv_holding is not None: + queues.append(kv_holding) return queues @@ -108,7 +106,7 @@ def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): return False if blocked_ids: - # A partner parked in skipped_waiting (e.g. waiting for remote KVs) + # A deferred partner (e.g. waiting for remote KVs) # may fail promotion this step while this member gets scheduled from # waiting; hold this member until the partner is schedulable. for rid in roles.values(): @@ -149,14 +147,14 @@ def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} scheduler._cfg_hold_counts = hold_counts - skipped = getattr(scheduler, "skipped_waiting", None) - blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() + kv_holding = getattr(scheduler, "kv_holding_waiting", None) + blocked_ids = {req.request_id for req in getattr(scheduler, "deferred_waiting", ())} held: list[tuple[Any, Any]] = [] for queue in _wait_queues(scheduler): # A member is only "blocked" for its partner's sake when it sits in # the OTHER queue; members of the queue being scanned move together. - partner_blockers = blocked_ids if queue is not skipped else set() + partner_blockers = blocked_ids if queue is not kv_holding else set() held_here = [] for request in list(queue): if _pair_complete(scheduler, request.request_id, partner_blockers): @@ -181,8 +179,12 @@ def _release_held(held: list[tuple[Any, Any]]) -> None: def _reorder_waiting_for_cfg(scheduler: Any) -> None: - """Move CFG pair partners adjacent in the FCFS waiting queue.""" - waiting = scheduler.waiting + """Move CFG pair partners adjacent in each FCFS waiting queue.""" + for waiting in _wait_queues(scheduler): + _reorder_cfg_queue(scheduler, waiting) + + +def _reorder_cfg_queue(scheduler: Any, waiting: Any) -> None: # Only the FCFS queue (a deque) has caller-controlled ordering; priority # queues order themselves and pair sync then relies on hold + equalize. if not isinstance(waiting, deque) or len(waiting) < 2: diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 09300af36..ba96348e0 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -28,7 +28,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -42,7 +41,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import random_sample from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.cosyvoice3.ras_sampler import MAX_FUSED_TOP_K, fused_ras_sample from vllm_omni.model_executor.models.cosyvoice3.runtime import ( cosyvoice3_batch_flow_debug, @@ -696,7 +695,7 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal ] -class CosyVoice3DummyInputsBuilder(BaseDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): +class CosyVoice3DummyInputsBuilder(OmniDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "Hello, this is a test of the CosyVoice3 system capability." diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index e97a78c25..ac45c0d5e 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py @@ -459,16 +459,22 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Process source images under vLLM's token-only prompt contract.""" mm_counts = mm_items.get_all_counts() num_images = mm_counts.get("image", 0) if num_images == 0: - return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_processor_mm_kwargs)}) + return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_kwargs)}) valid_mm_items = mm_items.select({key for key, count in mm_counts.items() if count > 0}) - processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + processor_data: dict[str, object] = {} + passthrough_data: dict[str, object] = {} + for items in valid_mm_items.values(): + if not items: + continue + processor_data.update(items.get_processor_data()) + passthrough_data.update(items.get_passthrough_data()) images = processor_data.get("images") if not images: return BatchFeature(dict(passthrough_data)) @@ -483,7 +489,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo source_grid_thw = image_grid_thw[:num_images] image_inputs["image_grid_thw"] = source_grid_thw - target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype) + target_grid = self._build_generation_grids(hf_kwargs)[:1].to(dtype=source_grid_thw.dtype) image_inputs["mrope_image_grid_thw"] = torch.cat( [source_grid_thw, target_grid], dim=0, diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index e98c227a4..773559339 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -38,7 +38,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -50,7 +49,7 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -674,7 +673,7 @@ class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessi ] -class GLMTTSDummyInputsBuilder(BaseDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]): +class GLMTTSDummyInputsBuilder(OmniDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "This is a test of the GLM-TTS voice cloning system." diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 4dbc9a455..338f145ad 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -29,7 +29,6 @@ from vllm.multimodal.inputs import ( ) from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, BaseProcessingInfo, ProcessorInputs, PromptReplacement, @@ -44,7 +43,7 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import ( NO_INTERLEAVE_NEXT_TOKEN_ID, @@ -224,7 +223,7 @@ class MiMoAudioLLMProcessingInfo( return {"audio": 1} -class MiMoAudioLLMDummyInputsBuilder(BaseDummyInputsBuilder[MiMoAudioLLMProcessingInfo]): +class MiMoAudioLLMDummyInputsBuilder(OmniDummyInputsBuilder[MiMoAudioLLMProcessingInfo]): _processor_inputs_cache: LRUCache = LRUCache(capacity=1024) def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: @@ -372,7 +371,7 @@ class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProces # Text-only input not supported in composite processor if not mm_data.get("audio", []): - prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids) + prompt_ids = self._postprocess_prompt(prompt_ids) return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index ec4faadb8..9e5517faa 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -396,25 +396,21 @@ class MimoAudioDummyInputsBuilder(BaseDummyInputsBuilder[MimoAudioProcessingInfo class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingInfo]): - def _preprocess_hf_mm_data( + def _get_hf_mm_inputs( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data( - mm_data, - hf_processor_mm_kwargs, - ) - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - if audios: - mm_data["audio"] = audios - feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) - hf_processor_mm_kwargs = { - **hf_processor_mm_kwargs, - "sampling_rate": feature_extractor.sampling_rate, - } - return mm_data, hf_processor_mm_kwargs + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ): + hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) + if hf_inputs.hf_data.get("audio"): + feature_extractor = self.info.get_feature_extractor(**hf_kwargs) + hf_inputs = hf_inputs._replace( + hf_kwargs={ + **hf_kwargs, + "sampling_rate": feature_extractor.sampling_rate, + }, + ) + return hf_inputs def _get_mm_fields_config( self, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 1338f2606..6887be33c 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -516,7 +516,7 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Call Ming's modality processors without expanding prompt text. @@ -525,7 +525,13 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO placeholder tokens remain under vLLM's prompt-update control. """ valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + mm_data: dict[str, object] = {} + passthrough_data: dict[str, object] = {} + for items in valid_mm_items.values(): + if not items: + continue + mm_data.update(items.get_processor_data()) + passthrough_data.update(items.get_passthrough_data()) hf_processor = self.info.get_hf_processor() data: dict[str, object] = {} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 2781fcc30..7f1d29345 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -3473,19 +3473,25 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """ Process each modality independently because the MiniCPM processor asserts that image tags and image sizes have matching lengths. """ valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + mm_data: dict[str, object] = {} + passthrough_data: dict[str, object] = {} + for items in valid_mm_items.values(): + if not items: + continue + mm_data.update(items.get_processor_data()) + passthrough_data.update(items.get_passthrough_data()) tokenizer = self.info.get_tokenizer() prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) input_ids = torch.tensor([tokenizer.encode(prompt_text)]) - mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs) + mm_inputs = self.process_mm_inputs(mm_data, hf_kwargs) processed_data = BatchFeature( { "input_ids": input_ids, diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py index 9ae74036e..f233287bb 100644 --- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -94,9 +94,9 @@ class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor): def _apply_hf_processor_main( self, mm_items: Any, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ): - base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs) + base_kwargs = self._base_processor_kwargs(hf_kwargs) return super()._apply_hf_processor_main( mm_items, base_kwargs, diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py index c382d4439..5a1c8f153 100644 --- a/vllm_omni/model_executor/models/moss_tts/local_model_state.py +++ b/vllm_omni/model_executor/models/moss_tts/local_model_state.py @@ -14,7 +14,10 @@ from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected +from vllm.v1.worker.gpu.sample.bad_words import BadWordsState +from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.penalties import PenaltiesState from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot @@ -118,15 +121,24 @@ class MossLocalModelState(OmniModelState): or sampler.return_sampling_mask or sampler.trace_replay_state is not None or sampler.get_logprobs_dims(rows) is not None - or np.any(sampler.logit_bias_state.use_logit_bias[rows]) - or np.any(sampler.penalties_state.use_penalty[rows]) - or np.any(sampler.bad_words_state.num_bad_words.np[rows]) or ( sampler.thinking_budget_state.enabled and np.any(sampler.thinking_budget_state.use_thinking_budget[rows]) ) ): return None + for processor in sampler.logits_processors: + if type(processor) is LogitBiasState: + active = np.any(processor.use_logit_bias[rows]) + elif type(processor) is PenaltiesState: + active = np.any(processor.use_penalty[rows]) + elif type(processor) is BadWordsState: + active = np.any(processor.num_bad_words.np[rows]) + else: + # Custom processors can alter even a single finite logit. + return None + if active: + return None keep = self.model._batch_should_continue if keep is None or keep.numel() != batch.num_reqs: return None diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index a0b0d725e..bca547eb8 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -28,7 +28,6 @@ from vllm.multimodal.inputs import ( ) from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -37,7 +36,7 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -117,7 +116,7 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Encode the reference audio, if any; the prompt was tokenized in apply(). @@ -126,7 +125,13 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr touch the text tokenizer. """ valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) - mm_data, passthrough = self._get_hf_mm_data(valid_items) + mm_data: dict[str, object] = {} + passthrough: dict[str, object] = {} + for items in valid_items.values(): + if not items: + continue + mm_data.update(items.get_processor_data()) + passthrough.update(items.get_passthrough_data()) audio = self._get_reference_audio(mm_data) if audio is None: return BatchFeature(dict(passthrough)) @@ -235,7 +240,7 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr ] -class OmniVoiceDummyInputsBuilder(BaseDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]): +class OmniVoiceDummyInputsBuilder(OmniDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "Hello, this is a test of the OmniVoice system." diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 1f89a37c4..8289a5fa2 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -78,7 +78,7 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, Vide from vllm.multimodal.processing.context import TimingContext from vllm.multimodal.processing.inputs import ProcessorInputs from vllm.multimodal.processing.processor import ( - MultiModalProcessingInfo, + MultiModalProcessingResult, MultiModalPromptUpdates, PlaceholderFeaturesInfo, PromptReplacement, @@ -367,7 +367,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], + hf_kwargs: Mapping[str, object], ) -> BatchFeature: """Apply omni-specific HF kwargs before the upstream Qwen2.5-Omni processor runs. @@ -378,23 +378,30 @@ class Qwen2_5OmniThinkerMultiModalProcessor( per-video ``use_audio_in_video`` mask handling have to run here. """ valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0}) - processor_data, _ = self._get_hf_mm_data(valid_mm_items) + # Upstream removed `_get_hf_mm_data`; extract the raw processor data + # (still keyed as `videos`/`video_metadata`/`audios`) for the omni + # kwargs helpers below, instead of `_get_hf_mm_inputs` which renames + # and injects a dummy text key. + processor_data: dict[str, object] = {} + for items in valid_mm_items.values(): + if items: + processor_data.update(items.get_processor_data()) - hf_processor_mm_kwargs = _presampled_videos_hf_kwargs( + hf_kwargs = _presampled_videos_hf_kwargs( processor_data, - hf_processor_mm_kwargs, + hf_kwargs, ) - hf_processor_mm_kwargs = _coerce_use_audio_in_video_for_hf_processor( + hf_kwargs = _coerce_use_audio_in_video_for_hf_processor( processor_data, - hf_processor_mm_kwargs, + hf_kwargs, ) processed_data = super()._apply_hf_processor_main( mm_items, - hf_processor_mm_kwargs, + hf_kwargs, ) - per_video_mask = hf_processor_mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) + per_video_mask = hf_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) if per_video_mask is not None: processed_data["use_audio_in_video"] = torch.tensor(per_video_mask) return processed_data @@ -564,32 +571,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ), ] - def _apply_hf_processor_mm_only( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ): - mm_counts = mm_items.get_all_counts() - - if "video" in mm_counts: - video_use_audio_in_video = _get_request_video_use_audio_in_video( - hf_processor_mm_kwargs, - mm_counts["video"], - ) - if any(video_use_audio_in_video): - assert "audio" in mm_counts - mm_counts["audio"] -= sum(video_use_audio_in_video) - - _, mm_processed_data, _ = self._apply_hf_processor_text_mm( - prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return mm_processed_data - def _get_audio_in_video_pairs( self, mm_data_items: MultiModalDataItems, @@ -639,14 +620,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> MultiModalProcessingInfo: + ) -> MultiModalProcessingResult: """Cache-aware: processes video/audio pairs as units, using pair-specific cache keys. Falls back to standard cache if no pairs.""" - cache = self.cache + cache = inputs.cache - _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items) - if cache is None or passthrough_data: + has_passthrough_data = any(len(items.get_passthrough_data()) > 0 for items in inputs.mm_data_items.values()) + if cache is None or has_passthrough_data: return self._apply_hf_processor(inputs, timing_ctx) aiv_pairs = self._get_audio_in_video_pairs( @@ -694,7 +675,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( with timing_ctx.record("apply_hf_processor"): mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, + hf_kwargs=hf_processor_mm_kwargs, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( @@ -720,7 +701,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor( mm_missing_prompt_updates=mm_missing_prompt_updates, ) - mm_info = MultiModalProcessingInfo( + mm_info = MultiModalProcessingResult( + prompt_ids=self._postprocess_prompt(inputs.prompt), kwargs=mm_kwargs, hashes=mm_hashes, prompt_updates=mm_prompt_updates, @@ -819,11 +801,12 @@ class Qwen2_5OmniThinkerMultiModalProcessor( def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, + mm_res: MultiModalProcessingResult, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() + prompt_ids = mm_res.prompt_ids + mm_kwargs = mm_res.kwargs + mm_prompt_updates = mm_res.prompt_updates self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index 6dca8b358..2122db233 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -561,35 +561,15 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing ) ] - def _preprocess_hf_mm_data( - self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], - ) -> tuple[Mapping[str, object], Mapping[str, object]]: - mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data( - mm_data, - hf_processor_mm_kwargs, - ) - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - if audios: - if isinstance(audios, str): - mm_data["audio"] = [audios] - elif isinstance(audios, (list, tuple)): - mm_data["audio"] = audios - else: - mm_data["audio"] = [audios] - return mm_data, hf_processor_mm_kwargs - def _postprocess_hf_mm_data( self, - mm_data: Mapping[str, object], - hf_processor_mm_kwargs: Mapping[str, object], + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], processed_data: BatchFeature, ) -> BatchFeature: processed_data = super()._postprocess_hf_mm_data( - mm_data, - hf_processor_mm_kwargs, + hf_data, + hf_kwargs, processed_data, ) if "audio_mels" not in processed_data: diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py index 05d73653f..0a4504820 100644 --- a/vllm_omni/model_executor/models/voxcpm2/scheduler.py +++ b/vllm_omni/model_executor/models/voxcpm2/scheduler.py @@ -30,7 +30,7 @@ class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler): def _should_defer_waiting_for_unified_decode_graph(self) -> bool: if not self._unified_decode_graph_enabled(): return False - if not self.waiting or not self.running: + if not (self.waiting or self.kv_holding_waiting) or not self.running: return False for request in self.running: diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 211159afa..f67edd30b 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -48,6 +48,8 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, Mult from vllm.multimodal.processing import BaseDummyInputsBuilder, BaseMultiModalProcessor from vllm.multimodal.processing.processor import ( BaseProcessingInfo, + MultiModalProcessingResult, + PlaceholderFeaturesInfo, ProcessorInputs, PromptReplacement, PromptUpdate, @@ -836,6 +838,14 @@ class VoxtralTTSDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralTTSProcessingIn class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessingInfo]): + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, Any], + ) -> ProcessorInputs: + return self.dummy_inputs.get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + def _get_mm_fields_config( self, hf_inputs: Mapping[str, NestedTensors], @@ -874,39 +884,17 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing ), ] - def _apply_hf_processor_mm_only( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Apply the HF processor on the multi-modal data only. - - Issue: Voxtral TTS use Mistral Tokenizer with custom audio encoder. It doesn't - inherit Transformers ProcessorMixin and can't use call_hf_processor_mm_only. - - Solution: Override this method to call _apply_hf_processor_text_mm directly. - """ - mm_counts = mm_items.get_all_counts() - _, mm_processed_data, _ = self._apply_hf_processor_text_mm( - prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - return mm_processed_data - def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates, - ): + mm_res: MultiModalProcessingResult, + ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: # Voxtral's Mistral chat template has already inserted the complete # audio-token run, so locate it without applying the replacement again. mm_item_counts = mm_items.get_all_counts() + prompt_ids = mm_res.prompt_ids + mm_kwargs = mm_res.kwargs + mm_prompt_updates = mm_res.prompt_updates self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) mm_placeholders = self._find_mm_placeholders( diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index ff0a21d45..791e91f5f 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -4,6 +4,7 @@ from dataclasses import dataclass, field from typing import Any +import numpy as np import torch from PIL import Image from vllm.outputs import CompletionOutput, RequestOutput @@ -86,6 +87,7 @@ _REQUEST_OUTPUT_CONTENT_ATTRS = ( "prompt", "prompt_token_ids", "prompt_logprobs", + "prompt_token_id_logprobs", "outputs", "finished", "lora_request", @@ -149,6 +151,7 @@ class OmniRequestOutput(RequestOutput): prompt: OmniPromptType | None = None prompt_token_ids: list[int] | None = None prompt_logprobs: Any = None + prompt_token_id_logprobs: np.ndarray | None = None outputs: list[CompletionOutput] = field(default_factory=list) finished: bool = True metrics: Any = field(default_factory=dict) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 85b583b99..d8ab8c2c6 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -10,6 +10,7 @@ from functools import cached_property import torch from aenum import extend_enum from vllm.config import ModelConfig as _OriginalModelConfig +from vllm.config import VllmConfig as _OriginalVllmConfig from vllm.inputs import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -29,6 +30,23 @@ from vllm_omni.request import OmniRequest, OmniStreamingUpdate _PATCH_LOGGER = logging.getLogger("vllm_omni.patch") +# Omni's deploy config selects a runner per stage. vLLM 0.31's config property +# otherwise defaults to v2 independently of the runner selected by our workers, +# making shared consumers such as InputProcessor choose the wrong API. +_runner_property = _OriginalVllmConfig.use_v2_model_runner +if not getattr(_runner_property.fget, "_omni_stage_runner", False): + + def _patched_use_v2_model_runner(self, _upstream=_runner_property.fget): + # Lazy import keeps the package's patch/config initialization acyclic. + from vllm_omni.config.model import OmniModelConfig + + if isinstance(self.model_config, OmniModelConfig): + return self.model_config.use_v2_model_runner + return _upstream(self) + + _patched_use_v2_model_runner._omni_stage_runner = True + _OriginalVllmConfig.use_v2_model_runner = property(_patched_use_v2_model_runner) + # ============================================================================= # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models # ============================================================================= @@ -390,8 +408,8 @@ def _patch_chat_template_registry(): ) if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: - _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = ( - lambda _: CHAT_TEMPLATES_DIR / "template_chatml.jinja" + _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = lambda _: ( + CHAT_TEMPLATES_DIR / "template_chatml.jinja" ) except ImportError: pass diff --git a/vllm_omni/platforms/npu/worker/aux_output.py b/vllm_omni/platforms/npu/worker/aux_output.py new file mode 100644 index 000000000..3fc093ae4 --- /dev/null +++ b/vllm_omni/platforms/npu/worker/aux_output.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Auxiliary-output delivery for Ascend's legacy model runner.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np + + +def _install_ascend_capture_sources(model: Any) -> None: + # AscendRoutedExperts invokes the router's private selection method, before + # logical-to-physical remapping and shared-expert expansion. It bypasses + # BaseRouter.select_experts(), where upstream normally invokes capture_fn. + for module in model.modules(): + if not type(module).__module__.startswith("vllm_ascend.") or not hasattr(module, "layer_id"): + continue + router = getattr(module, "router", None) + select_experts = getattr(router, "_select_experts", None) + if not callable(select_experts) or getattr(module, "_omni_npu_capture_source", False): + continue + module.capture_fn = None + + def select_and_capture(*args, _select=select_experts, _source=module, **kwargs): + weights, logical_ids = _select(*args, **kwargs) + if _source.capture_fn is not None: + _source.capture_fn(logical_ids) + return weights, logical_ids + + router._select_experts = select_and_capture + # layer_id + capture_fn implement upstream RoutedExpertsCaptureSource. + module._omni_npu_capture_source = True + + +class NPUAuxOutputMixin: + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.aux_output_connector = None + + def _init_omni_aux_output(self, kv_cache_config) -> None: + if not self.vllm_config.aux_output_config.enabled: + return + from vllm.distributed.aux_output_connector.worker import get_aux_output_connector + + self._close_omni_aux_output() + _install_ascend_capture_sources(self.model) + self.aux_output_connector = get_aux_output_connector(self.model, self.vllm_config, kv_cache_config) + + def _begin_omni_aux_output_step(self, scheduler_output) -> None: + if self.aux_output_connector is not None: + # Apply terminal/hash-only metadata even on no-forward steps. + self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata) + + def _prepare_omni_aux_output(self): + if self.aux_output_connector is None: + return None + num_reqs = len(self.input_batch.req_ids) + # Bookkeeping and the next async batch mutate these legacy buffers. + batch = SimpleNamespace( + req_ids=list(self.input_batch.req_ids), + query_start_loc_np=self.query_start_loc.np[: num_reqs + 1].copy(), + num_computed_tokens_np=self.input_batch.num_computed_tokens_cpu[:num_reqs].copy(), + ) + return self.aux_output_connector.prepare_output(batch) + + def _finish_omni_aux_output(self, pending, scheduler_output, sampled_token_ids=None, invalid_req_indices=()): + if pending is None: + return None + num_reqs = len(pending.request_ids) + num_sampled = np.zeros(num_reqs, dtype=np.int32) + num_rejected = np.zeros(num_reqs, dtype=np.int32) + if sampled_token_ids is not None: + # Instrumentation opts into a synchronized copy. The normal NPU + # async sampling path remains unchanged when capture is disabled. + sampled = sampled_token_ids.detach().to("cpu").numpy() + num_sampled[:] = (sampled >= 0).sum(axis=-1) + for i, request_id in enumerate(pending.request_ids): + draft_count = len(scheduler_output.scheduled_spec_decode_tokens.get(request_id, ())) + if draft_count: + num_rejected[i] = max(0, draft_count + 1 - int(num_sampled[i])) + # Partial prefills execute valid routing rows but emit no token. + num_sampled[list(invalid_req_indices)] = 0 + pending.enqueue_cpu_copy(num_sampled=num_sampled, num_rejected=num_rejected) + self._sync_device() + return pending.process_output() + + def _close_omni_aux_output(self) -> None: + connector = self.aux_output_connector + self.aux_output_connector = None + if connector is not None: + try: + connector.close() + finally: + for module in self.model.modules(): + if getattr(module, "_omni_npu_capture_source", False): + module.capture_fn = None diff --git a/vllm_omni/platforms/npu/worker/base.py b/vllm_omni/platforms/npu/worker/base.py index 5dd929ea1..a79290640 100644 --- a/vllm_omni/platforms/npu/worker/base.py +++ b/vllm_omni/platforms/npu/worker/base.py @@ -59,3 +59,11 @@ class OmniNPUWorkerBase(NPUWorker): self.profiler.start() else: self.profiler.stop() + + def shutdown(self) -> None: + try: + runner = getattr(self, "model_runner", None) + if runner is not None: + runner._close_omni_aux_output() + finally: + super().shutdown() diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index ee5407065..082b087fa 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -315,10 +315,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: - if self.vllm_config.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): @@ -392,6 +388,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex get_kv_transfer_group().handle_preemptions(kv_connector_metadata) # -------------------------------------- Omni-new ------------------------------------------------- + self._begin_omni_aux_output_step(scheduler_output) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with record_function_or_nullcontext("prepare input"): with self.synchronize_input_prep(): @@ -791,9 +788,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None def _sample( @@ -884,6 +878,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex multimodal_outputs, # Omni-Specific prefix_cache_step_id, ) = self.execute_model_state + pending_aux_output = self._prepare_omni_aux_output() # Clear ephemeral state. self.execute_model_state = None hidden_seq_len = int(hidden_states.shape[0]) @@ -1015,11 +1010,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex draft_by_req_id = dict(zip(draft_req_ids, draft_ids_list)) output_spec_token_ids = [draft_by_req_id.get(req_id, []) for req_id in req_ids_output_copy] - routed_experts_lists = None - if self.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"): - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) + aux_output = self._finish_omni_aux_output( + pending_aux_output, scheduler_output, sampler_output.sampled_token_ids, invalid_req_indices + ) # -------------------------------------- Omni-new ------------------------------------------------- engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) @@ -1209,7 +1202,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex num_nans_in_logits=num_nans_in_logits, ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids - model_runner_output.routed_experts = routed_experts_lists + model_runner_output.aux_output_connector_output = aux_output model_runner_output.spec_token_ids = output_spec_token_ids with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"): model_runner_output.omni_connector_output = self.get_omni_connector_output() diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index bc486e4b5..8948f3125 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -89,10 +89,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: - if self.vllm_config.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): @@ -127,6 +123,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if flush_ids: self.flush_full_payload_outputs(flush_ids) + self._begin_omni_aux_output_step(scheduler_output) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with record_function_or_nullcontext("prepare input"): # -------------------------------------- Omni-new ------------------------------------------------- @@ -418,9 +415,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None @torch.inference_mode() @@ -461,6 +455,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin multimodal_outputs_raw, # Omni-Specific _prefix_cache_step_id, # generation stages never save to the prefix cache ) = self.execute_model_state + pending_aux_output = self._prepare_omni_aux_output() # Clear ephemeral state. self.execute_model_state = None @@ -518,11 +513,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin req_state = self.requests.get(rid) if req_state is not None and inter_stage_outputs[i]: self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state) - routed_experts_lists = None - if self.vllm_config.model_config.enable_return_routed_experts and hasattr( - self.input_batch, "num_tokens_no_spec" - ): - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) + aux_output = self._finish_omni_aux_output(pending_aux_output, scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -537,7 +528,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_stats=cudagraph_stats, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) - output.routed_experts = routed_experts_lists + output.aux_output_connector_output = aux_output output.omni_connector_output = self.get_omni_connector_output() # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index da663f626..a490950a5 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -24,6 +24,7 @@ from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE from vllm_omni.core.prefix_cache import stage_prefix_cache_config from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p +from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) @@ -35,7 +36,7 @@ else: from vllm_ascend.worker.model_runner_v1 import NPUModelRunner -class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): +class OmniNPUModelRunner(NPUAuxOutputMixin, OmniGPUModelRunner, NPUModelRunner): def initialize_kv_cache( self, kv_cache_config, kv_cache_allocation_context: AbstractContextManager | None = None ) -> None: @@ -43,6 +44,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): NPUModelRunner.initialize_kv_cache( self, kv_cache_config, kv_cache_allocation_context=kv_cache_allocation_context ) + self._init_omni_aux_output(kv_cache_config) if getattr(self, "_omni_prefix_cache_cfg", None) is None: # Same gate as the GPU runner (pooling stage, kv_consumer / # kv_both, hybrid kv groups). Read the config back off diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index c0975f651..55a50a090 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -81,6 +81,16 @@ from .component_config import ComponentQuantizationConfig # noqa: E402 logger = init_logger(__name__) +def _build_fp8(**kw: Any) -> QuantizationConfig: + if kw.get("is_checkpoint_fp8_serialized", False): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + return Fp8Config(**kw) + from .fp8_config import DiffusionFp8Config + + return DiffusionFp8Config(**kw) + + def _build_int8(**kw: Any) -> QuantizationConfig: """Lazy import for Int8 diffusion config (supports CUDA + NPU).""" from .int8_config import DiffusionInt8Config @@ -168,6 +178,7 @@ def _build_torchao_float8_weight_only(**kw: Any) -> QuantizationConfig: _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { + "fp8": _build_fp8, "int8": _build_int8, "bitsandbytes": _build_bitsandbytes, "mxfp8": _build_mxfp8, @@ -330,6 +341,17 @@ def _build_single(method: str, **kwargs: Any) -> QuantizationConfig: config_cls = get_quantization_config(method) + from vllm.config.quantization import resolve_quantization_config + from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + + if config_cls is OnlineQuantizationConfig: + if "args" in kwargs: + return config_cls(**kwargs) + args = resolve_quantization_config(method, kwargs or None) + if args is None: + raise ValueError("Online quantization requires quantization config arguments") + return config_cls(args) + try: return config_cls(**kwargs) except TypeError: diff --git a/vllm_omni/quantization/fp8_config.py b/vllm_omni/quantization/fp8_config.py new file mode 100644 index 000000000..4f610d8e9 --- /dev/null +++ b/vllm_omni/quantization/fp8_config.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni's online FP8 option backed by vLLM's online quantization API.""" + +from vllm.config.quantization import QuantizationConfigArgs, QuantSpec +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + + +class DiffusionFp8Config(OnlineQuantizationConfig): + """Keep the diffusion ``fp8`` option for quantizing BF16/FP16 weights. + + Upstream Fp8Config now describes serialized checkpoints only. Keeping the + Omni method name also lets checkpoint metadata replace this online config + through the existing serialized-checkpoint reconciliation path. + """ + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + weight_block_size: list[int] | None = None, + store_dtype: str | None = None, + ) -> None: + if is_checkpoint_fp8_serialized: + raise ValueError("Serialized FP8 checkpoints require upstream Fp8Config") + if activation_scheme != "dynamic": + raise ValueError("Online FP8 quantization requires activation_scheme='dynamic'") + if weight_block_size is not None: + raise ValueError("Block-wise FP8 quantization requires a serialized checkpoint or fp8_per_block") + if store_dtype is not None: + raise ValueError("Online FP8 quantization does not support store_dtype") + super().__init__( + QuantizationConfigArgs( + linear=QuantSpec(weight="fp8_per_tensor_static"), + moe=QuantSpec(weight="fp8_per_tensor_static"), + ignore=list(ignored_layers or []), + ) + ) + self.is_checkpoint_fp8_serialized = False + self.activation_scheme = activation_scheme + self.weight_block_size = None + self.store_dtype = None + + @classmethod + def get_name(cls) -> str: + return "fp8" diff --git a/vllm_omni/request.py b/vllm_omni/request.py index a6da16696..643ab020a 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -107,6 +107,7 @@ class OmniRequest(Request): model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), resumable=request.resumable, session_id=request.session_id, + kv_hints=request.kv_hints, 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 0a9cdd7b0..78e653b54 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -206,7 +206,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): invalid_req_indices = kwargs.pop("invalid_req_indices") 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 @@ -223,7 +222,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): self._sampled_token_ids = sampled_token_ids self.vocab_size = vocab_size 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 @@ -240,9 +238,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): # output asynchronously. self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu", non_blocking=True) self._logprobs_tensors_cpu = self._logprobs_tensors.to_cpu_nonblocking() if self._logprobs_tensors else None - 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() @@ -940,9 +935,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - if self.routed_experts_initialized: - self.routed_experts_capturer.clear_buffer() - if not getattr(self, "_warmup_state_cleared", False): self._warmup_state_cleared = True if hasattr(self.model, "_clear_warmup_state"): @@ -974,11 +966,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if flush_ids: self.flush_full_payload_outputs(flush_ids) - if self.routed_experts_initialized: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() - # If ngram_gpu is used, we need to copy the scheduler_output to avoid # the modification has influence on the scheduler_output in engine core process. # The replace is much faster than deepcopy. @@ -1371,9 +1358,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self._should_return_omni_routed_experts() and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None def _sample( @@ -1628,14 +1612,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex except TypeError: return scheduler_output - def _should_return_omni_routed_experts(self) -> bool: - model_config = getattr(self, "model_config", None) - if model_config is None: - model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) - return bool(getattr(model_config, "enable_return_routed_experts", False)) and bool( - getattr(self, "routed_experts_initialized", False) - ) - @staticmethod def _model_omni_flag(model: Any, name: str, default: bool = False) -> bool: return bool(getattr(model, name, default)) if model is not None else default @@ -1700,8 +1676,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex getattr(self, "model", None), "supports_async_whole_payload" ): return False - if bool(getattr(model_config, "enable_return_routed_experts", False)): - return False model = getattr(self, "model", None) if not self._model_omni_flag(model, "use_async_omni_output"): @@ -1997,9 +1971,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): - routed_experts_lists = None - if self._should_return_omni_routed_experts(): - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -2015,7 +1986,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats=cudagraph_stats, ) output.kv_extracted_req_ids = kv_extracted_req_ids - output.routed_experts = routed_experts_lists return output @torch.inference_mode() diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index a05e12708..7c2cf2163 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -137,9 +137,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - if self.routed_experts_initialized: - self.routed_experts_capturer.clear_buffer() - if hasattr(self, "_omni_connector"): for request in getattr(scheduler_output, "pending_input_registrations", []): self.register_chunk_recv(request) @@ -150,11 +147,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if flush_ids: self.flush_full_payload_outputs(flush_ids) - if self.routed_experts_initialized: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() - # If ngram_gpu is used, we need to copy the scheduler_output to avoid # the modification has influence on the scheduler_output in engine core process. # The replace is much faster than deepcopy. @@ -401,9 +393,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.routed_experts_initialized and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None @torch.inference_mode() @@ -516,9 +505,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() - routed_experts_lists = None - if self.routed_experts_initialized: - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) if inter_stage_outputs and self._should_accumulate_full_payload_output(): for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) @@ -541,7 +527,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) output.omni_connector_output = self.get_omni_connector_output() - output.routed_experts = routed_experts_lists if not self.use_async_scheduling: return output @@ -884,7 +869,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # in code2wav warmup. The expert-balance rationale behind # randomize_inputs targets text-MoE models, which generation # stages are not. The kwarg stays accepted for vLLM API compat. - self.maybe_randomize_inputs(input_ids, inputs_embeds), set_forward_context( attn_metadata, self.vllm_config, @@ -964,7 +948,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin return hidden_states, None - def profile_run(self) -> None: + def profile_run(self, randomize_inputs: bool = False) -> None: # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config @@ -1017,7 +1001,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self.encoder_cache[f"tmp_{i}"] = output # Add `is_profile` here to pre-allocate communication buffers - hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) + hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True, randomize_inputs=randomize_inputs) output = None self._sync_device() del hidden_states diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d8af6b201..931e36d1c 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -45,7 +45,6 @@ from vllm_omni.utils.device_copy import index_to_device if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput - from vllm.v1.outputs import RoutedExpertsLists else: xgr = LazyLoader("xgr", globals(), "xgrammar") xgr_torch_compile = LazyLoader( @@ -108,52 +107,6 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): return sampled return super()._to_list(sampled_token_ids) - def _omni_routed_experts_d2h(self, scheduler_output) -> None: - """Issue routed-experts D2H copy matching upstream GPUModelRunner pattern. - - Upstream does this inline in ``execute_model``: - buf = self.routed_experts_capturer.get_device_buffer() - total = scheduler_output.total_num_scheduled_tokens - self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) - self.routed_experts_slot_mapping_cpu[:total].copy_( - self.routed_experts_slot_mapping_device[:total], non_blocking=True) - """ - if not self.routed_experts_initialized: - return - buf = self.routed_experts_capturer.get_device_buffer() - total = scheduler_output.total_num_scheduled_tokens - self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) - if hasattr(self, "routed_experts_slot_mapping_device"): - self.routed_experts_slot_mapping_cpu[:total].copy_( - self.routed_experts_slot_mapping_device[:total], - non_blocking=True, - ) - - def _omni_extract_routed_experts(self, scheduler_output) -> "RoutedExpertsLists | None": - """Extract routed experts matching upstream GPUModelRunner pattern. - - Upstream (sync path, sample_tokens): - total = scheduler_output.total_num_scheduled_tokens - output.routed_experts = RoutedExpertsLists( - routing_data=self.routed_experts_cpu[:total].numpy(), - slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), - ) - - Returns RoutedExpertsLists (batch-level, with slot_mapping) so that - downstream schedulers can use slot_mapping to map back to requests. - """ - from vllm.v1.outputs import RoutedExpertsLists - - if not self.routed_experts_initialized: - return None - total = scheduler_output.total_num_scheduled_tokens - if total <= 0: - return None - return RoutedExpertsLists( - routing_data=self.routed_experts_cpu[:total].numpy(), - slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), - ) - def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes): """Initialize metadata builders and keep FA3 graph metadata buffers sized. diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py index 7475d3d7d..50386df02 100644 --- a/vllm_omni/worker_v2/model_states/__init__.py +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -4,7 +4,8 @@ """Omni-aware ``init_model_state`` factory. Extends the upstream v2 factory with Omni architecture dispatch. -Non-Omni architectures fall through to the upstream ``init_model_state``. +Non-Omni architectures use upstream dispatch; its default state acquires +Omni lifecycle support when owned by an Omni stage. """ from __future__ import annotations @@ -16,8 +17,11 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states import ( init_model_state as _upstream_init_model_state, ) +from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm_omni.config.model import OmniModelConfig + def init_omni_model_state( vllm_config: VllmConfig, @@ -28,8 +32,8 @@ def init_omni_model_state( """Create the appropriate ``ModelState`` for *model*. Returns an ``OmniModelState`` when the model declares Omni lifecycle - capabilities; otherwise - delegates to the upstream v2 factory. + capabilities. Otherwise delegates to the upstream v2 factory, wrapping + its default state for auxiliary Omni stages. """ factory = getattr(model, "create_omni_model_state", None) if factory is not None: @@ -49,4 +53,12 @@ def init_omni_model_state( ) return OmniModelState(vllm_config, model, encoder_cache, device) - return _upstream_init_model_state(vllm_config, model, encoder_cache, device) + state = _upstream_init_model_state(vllm_config, model, encoder_cache, device) + if type(state) is DefaultModelState and isinstance(getattr(vllm_config, "model_config", None), OmniModelConfig): + # Auxiliary stages such as the forced aligner have no model-owned + # hooks, but Omni's runner still needs its request/output lifecycle. + # Keep specialized upstream states and plain vLLM configs intact. + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + + return OmniModelState(vllm_config, model, encoder_cache, device) + return state diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 80eb56b83..57bbf9b5b 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -31,6 +31,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.utils import AttentionGroup @@ -447,6 +448,7 @@ class OmniModelState(DefaultModelState): kv_cache_config: KVCacheConfig, for_capture: bool = False, ubatch_idx: int = 0, + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, ) -> dict[str, Any]: if ( for_capture @@ -471,6 +473,7 @@ class OmniModelState(DefaultModelState): kv_cache_config, for_capture=for_capture, ubatch_idx=ubatch_idx, + model_specific_attn_metadata=model_specific_attn_metadata, ) def prepare_inputs_embeds( diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index c9ef50236..11b14e2d9 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -25,7 +25,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import ( ) from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput, RoutedExpertsTensors +from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput from vllm.v1.worker.gpu.eplb_utils import step_eplb_after from vllm_omni.data_entry_keys import OmniPayload, flatten_payload @@ -121,6 +121,7 @@ class OmniARModelRunner(OmniGPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> Any: if not dummy_run: self._handle_kv_transfer_pre(scheduler_output) @@ -132,6 +133,7 @@ class OmniARModelRunner(OmniGPUModelRunner): is_profile=is_profile, context_len=context_len, valid_dummy_state_slots=valid_dummy_state_slots, + randomize_inputs=randomize_inputs, ) # ------------------------------------------------------------------ @@ -170,7 +172,9 @@ class OmniARModelRunner(OmniGPUModelRunner): hidden_states = self.execute_model_state.hidden_states finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output - routed_experts = self.execute_model_state.routed_experts + pending_aux_output = None + if self.aux_output_connector is not None: + pending_aux_output = self.aux_output_connector.prepare_output(input_batch) self.execute_model_state = None # --- Omni: reconstruct raw model output and post-process --- @@ -253,6 +257,12 @@ class OmniARModelRunner(OmniGPUModelRunner): self.req_states.num_computed_tokens.gpu, self.req_states.prompt_len.np, ) + prompt_token_id_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_token_id_logprobs( + self.model.compute_logits, + text_hidden, + input_batch, + self.req_states.prompt_len.np, + ) # --- Omni: pooler_output --- engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", "text") @@ -264,6 +274,7 @@ class OmniARModelRunner(OmniGPUModelRunner): req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)}, sampled_token_ids=None, # type: ignore[arg-type] prompt_logprobs_dict=prompt_logprobs_dict, + prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict, kv_connector_output=None, ) model_runner_output.kv_extracted_req_ids = kv_extracted @@ -291,7 +302,7 @@ class OmniARModelRunner(OmniGPUModelRunner): finalize_output=(None if materialize_native else self._finalize_native_data_plane_output), finalize_multimodal=sampling_output.finalize_multimodal, check_ep_fault=self.check_ep_fault, - routed_experts=routed_experts, + pending_aux_output=pending_aux_output, streaming_audio=streaming_audio, extra_multimodal_outputs=extra_outputs, ) @@ -722,14 +733,14 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): finalize_output: Any | None = None, finalize_multimodal: Any | None = None, check_ep_fault: bool = False, - routed_experts: RoutedExpertsTensors | None = None, + pending_aux_output: Any | None = None, streaming_audio: StreamingAudioOutput | None = None, extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None, ): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.routed_experts = routed_experts + self.pending_aux_output = pending_aux_output self.copy_event = copy_event if copy_event is not None else torch.cuda.Event(blocking=True) self._async_chunk = bool(async_chunk) self._finalize_output = finalize_output @@ -784,9 +795,15 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self.sampling_mask_tensors = None if sampler_output.sampling_mask_tensors is not None: self.sampling_mask_tensors = sampler_output.sampling_mask_tensors.to_cpu_nonblocking() - self.routed_experts_cpu = None - if routed_experts is not None: - self.routed_experts_cpu = routed_experts.to_cpu_nonblocking() + if pending_aux_output is not None: + pending_aux_output.enqueue_cpu_copy( + num_sampled=self.num_sampled_tokens_np, + num_rejected=_async_copy_to_np( + sampler_output.num_rejected, + copy_stream=copy_stream, + pin_memory=pin_memory, + ), + ) # Logprobs self.logprobs_tensors = None @@ -805,6 +822,10 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): k: v.to_cpu_nonblocking() if v is not None else None for k, v in self.model_runner_output.prompt_logprobs_dict.items() } + self.prompt_token_id_logprobs_dict = { + k: _async_copy_tensor(v, copy_stream=copy_stream, pin_memory=pin_memory) + for k, v in self.model_runner_output.prompt_token_id_logprobs_dict.items() + } if check_ep_fault: has_fault = get_ep_all2all_manager().query_fault() self._has_fault = has_fault.to("cpu", non_blocking=True) @@ -879,8 +900,8 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self.model_runner_output.sampled_token_ids = sampled_token_ids if self.sampling_mask_tensors is not None: self.model_runner_output.sampling_masks = self.sampling_mask_tensors.tolists() - if self.routed_experts_cpu is not None: - self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists() + if self.pending_aux_output is not None: + self.model_runner_output.aux_output_connector_output = self.pending_aux_output.process_output() self.model_runner_output.sampled_token_ids_materialized = True # Logprobs @@ -891,6 +912,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): if self.logprobs_tensors is not None: self.model_runner_output.logprobs = self.logprobs_tensors.tolists() self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict + self.model_runner_output.prompt_token_id_logprobs_dict = self.prompt_token_id_logprobs_dict if self._has_fault is not None and self._has_fault.item(): mask = get_ep_all2all_manager().query_active_mask() diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py index 605531338..abde9a7db 100644 --- a/vllm_omni/worker_v2/omni_generation_model_runner.py +++ b/vllm_omni/worker_v2/omni_generation_model_runner.py @@ -14,6 +14,7 @@ from __future__ import annotations from dataclasses import replace from typing import Any +import numpy as np import torch from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import set_forward_context @@ -123,10 +124,12 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput): copy_stream: torch.cuda.Stream, finalize_output: Any | None = None, check_ep_fault: bool = False, + pending_aux_output: Any | None = None, ) -> None: self.model_runner_output = model_runner_output self.num_reqs = num_reqs self.finalize_output = finalize_output + self.pending_aux_output = pending_aux_output self.copy_event = torch.cuda.Event(blocking=True) self._has_fault: torch.Tensor | None = None @@ -138,6 +141,10 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput): copy_stream=copy_stream, pin_memory=PIN_MEMORY, ) + if pending_aux_output is not None: + # Generation stages emit no sampled/rejected token IDs. + counts = np.zeros(num_reqs, dtype=np.int32) + pending_aux_output.enqueue_cpu_copy(num_sampled=counts, num_rejected=counts) if check_ep_fault: has_fault = get_ep_all2all_manager().query_fault() self._has_fault = has_fault.to("cpu", non_blocking=True) @@ -145,6 +152,8 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput): def get_output(self) -> OmniModelRunnerOutput: self.copy_event.synchronize() + if self.pending_aux_output is not None: + self.model_runner_output.aux_output_connector_output = self.pending_aux_output.process_output() if self._has_fault is not None and self._has_fault.item(): mask = get_ep_all2all_manager().query_active_mask() raise RuntimeError( @@ -283,7 +292,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): # ------------------------------------------------------------------ @torch.inference_mode() - def profile_run(self) -> None: + def profile_run(self, randomize_inputs: bool = False) -> None: """Generation models have no KV cache — skip profiling. Code2Wav shares GPU memory with the Talker stage (same device); @@ -308,6 +317,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: self._prepare_native_data_plane(scheduler_output) @@ -321,6 +331,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): # Registration/terminal/abort work has been applied above. # There is no new or cached request state to add or update. self._apply_block_table_staged_writes_if_available() + if self.aux_output_connector is not None: + self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata) empty_output = self.kv_connector.no_forward(scheduler_output) return self._attach_native_data_plane_signals( self._merge_ec_connector_no_forward(scheduler_output, empty_output) @@ -333,6 +345,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): self.update_requests(scheduler_output) self._sync_native_data_plane_payloads(scheduler_output) self._apply_block_table_staged_writes_if_available() + if self.aux_output_connector is not None: + self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata) if scheduler_output.total_num_scheduled_tokens == 0: empty_output = self.kv_connector.no_forward(scheduler_output) return self._attach_native_data_plane_signals( @@ -366,7 +380,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): if not dummy_run: assert batch_req_state is not None - input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc) + input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc, 0) else: from vllm.v1.worker.gpu.input_batch import InputBatch @@ -375,7 +389,10 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + is_padding=not is_profile, ) + # Codec IDs have model-specific bounds; keep valid dummy tokens + # even when upstream requests randomized language-model inputs. attn_metadata = None slot_mappings_by_layer = None @@ -451,7 +468,6 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): finished_req_ids=scheduler_output.finished_req_ids, dp_sync=dp_sync, ec_connector_output=ec_connector_output, - routed_experts=None, cudagraph_stats=None, ) return None @@ -476,6 +492,9 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): return None num_reqs = input_batch.num_reqs + pending_aux_output = None + if self.aux_output_connector is not None: + pending_aux_output = self.aux_output_connector.prepare_output(input_batch) # Mark all scheduled tokens as computed so the scheduler does # not re-schedule them. Unlike AR stages we do NOT call @@ -517,7 +536,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): raw_multimodal_outputs = model_output.multimodal_outputs can_copy_async = getattr(getattr(self, "device", None), "type", "cpu") == "cuda" and ( - self.check_ep_fault or _contains_cuda_tensor(raw_multimodal_outputs) + self.check_ep_fault or pending_aux_output is not None or _contains_cuda_tensor(raw_multimodal_outputs) ) if can_copy_async: async_output = OmniGenerationAsyncOutput( @@ -528,6 +547,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): copy_stream=self.output_copy_stream, finalize_output=self._finalize_native_data_plane_output, check_ep_fault=self.check_ep_fault, + pending_aux_output=pending_aux_output, ) self._reserve_native_data_plane_outputs(list(req_ids)) self._release_generation_slots(input_batch) diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 9e73201e5..c07c47f89 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -391,7 +391,7 @@ class OmniGPUModelRunner(GPUModelRunner): } logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.") - def capture_model(self) -> int: + def capture_model(self, *, profile_only: bool = False) -> int: """Handle CUDA graph capture for Omni models. Tuple-returning models use PIECEWISE graphs; FULL replay requires @@ -424,12 +424,17 @@ class OmniGPUModelRunner(GPUModelRunner): use_aux = self.use_aux_hidden_state_outputs self.use_aux_hidden_state_outputs = use_aux or aux_outputs try: - result = super().capture_model() + result = super().capture_model(profile_only=profile_only) finally: self.model.forward = original_forward # type: ignore[assignment] self.use_aux_hidden_state_outputs = use_aux else: - result = super().capture_model() + result = super().capture_model(profile_only=profile_only) + + if profile_only: + # Upstream discards captures in its temporary profiling pool. + # Model-owned graphs must be recorded during the real capture. + return result capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None) if callable(capture_mtp): @@ -494,6 +499,7 @@ class OmniGPUModelRunner(GPUModelRunner): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> Any: if not dummy_run: self._prepare_native_data_plane(scheduler_output) @@ -503,6 +509,8 @@ class OmniGPUModelRunner(GPUModelRunner): self.update_requests(scheduler_output) self._sync_native_data_plane_payloads(scheduler_output) self.block_tables.apply_staged_writes() + if self.aux_output_connector is not None: + self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata) if scheduler_output.total_num_scheduled_tokens == 0: empty_output = self.kv_connector.no_forward(scheduler_output) return self._attach_native_data_plane_signals( @@ -544,7 +552,7 @@ class OmniGPUModelRunner(GPUModelRunner): if not dummy_run: assert batch_req_state is not None - input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc) + input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc, num_active_loras) block_tables, slot_mappings = self.prepare_attn(input_batch) self.model_state.preprocess_state( input_batch, @@ -568,7 +576,10 @@ class OmniGPUModelRunner(GPUModelRunner): batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + is_padding=not is_profile, ) + if randomize_inputs: + input_batch.input_ids.random_(0, self.vocab_size) if not skip_attn_for_dummy_run: block_tables, slot_mappings = self.prepare_dummy_attn(input_batch, valid_dummy_state_slots) if context_len: @@ -685,11 +696,6 @@ class OmniGPUModelRunner(GPUModelRunner): if not dummy_run and isinstance(hidden_states, torch.Tensor): self.model_state.run_postprocess(hidden_states, input_batch) - routed_experts = None - if not dummy_run and (capturer := self.routed_experts_capturer) is not None: - assert slot_mappings is not None - routed_experts = capturer.get_routed_experts(slot_mappings, num_toks) - self.execute_model_state = ExecuteModelState( input_batch=input_batch, attn_metadata=attn_metadata, @@ -699,7 +705,6 @@ class OmniGPUModelRunner(GPUModelRunner): finished_req_ids=scheduler_output.finished_req_ids, dp_sync=dp_sync, ec_connector_output=ec_connector_output, - routed_experts=routed_experts, cudagraph_stats=None, )