diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index d638ecb8d..f51f13f1a 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -28,6 +28,9 @@ # packed onto g6.12xlarge L4 nodes (``vllm.ci/gpu-pool=l4x4``). GPU / CPU / # memory / shm requests scale with card count (1 GPU = 10 CPU / 40Gi / 24Gi shm). # Retry covers lost agents / killed pods, not pytest or OOM (exit 1 / 137). +# The Kubernetes HF secret deliberately uses VLLM_CI_HF_TOKEN. Scheduled +# build-level HF_TOKEN values can override pod env entries with the same name; +# upload_pipeline.py restores the secret-backed value inside the job shell. x-npu-k8s-plugins: &npu_k8s_plugins - kubernetes: @@ -86,7 +89,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -137,7 +140,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -188,7 +191,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -239,7 +242,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -273,7 +276,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -308,7 +311,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -343,7 +346,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -378,7 +381,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -428,7 +431,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -479,7 +482,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -530,7 +533,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -581,7 +584,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 046b136d7..749f35f5a 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -87,6 +87,8 @@ BOOTSTRAP_UPLOAD_IF_KEYS = { } E2E_GROUP_MARKER = "E2E Test" CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml" +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' # Bootstrap Buildkite ``if`` expressions. # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN). @@ -509,6 +511,17 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: expanded = copy.deepcopy(preset) merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded + # A scheduled build's HF_TOKEN may override a pod env entry with the same + # name. The preset injects the Kubernetes secret under a collision-proof + # alias; restore the conventional name after Buildkite has merged env. + 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] + elif commands is None: + merged["commands"] = [CUDA_HF_TOKEN_EXPORT] + else: + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, 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/docker/Dockerfile.ci b/docker/Dockerfile.ci index b84b344a8..a8e46580b 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,27 +1,10 @@ -# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5). -# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. -# -# Release-boundary form: there is no precompiled-wheel pin here, because the -# base image already ships exactly the target vLLM. Everything the runtime -# needs — torch, flashinfer, nvJitLink, numpy — comes from that image as one -# self-consistent set. -# -# This block reappears only when the rebase targets a vLLM commit that is NOT a -# released tag. During the 0.27 rebase, targeting v0.27.0rc1 on top of the -# v0.26.0 base image required reinstalling vLLM from wheels.vllm.ai and then -# hand-repairing the fallout: pinning torch ahead of flashinfer (installing it -# after bound flashinfer's TVM FFI runtime to the old torch and aborted engine -# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy -# under numba's ceiling. All of it was compensation for a base image and a vLLM -# that disagreed. At a real release boundary they agree, so none of it is -# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels -# rode on the v0.27.0 base image with the full reinstall block until the -# official v0.28.0 image shipped, at which point this file dropped back to the -# form you are reading. If you pin an untagged commit again, restore the -# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather -# than letting the Dockerfile silently ship the wrong vLLM. +# Target vLLM version: v0.29.0 (tag @ 98dff2a81d74). +# 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. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.28.0 +ARG VLLM_BASE_TAG=v0.29.0 # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents diff --git a/docs/configuration/README.md b/docs/configuration/README.md index cdfb5818a..3d065f9a1 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.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html). +For options within a vLLM Engine, please refer to the [vLLM 0.29 configuration guide](https://docs.vllm.ai/en/v0.29.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 8f944f44d..90b9ffae0 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -151,7 +151,7 @@ their keys only. ## Inherited vLLM variables vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to -the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html) +the [vLLM 0.29 environment-variable reference](https://docs.vllm.ai/en/v0.29.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 448ca81f2..9d2d11e94 100644 --- a/docs/getting_started/installation/README.md +++ b/docs/getting_started/installation/README.md @@ -1,7 +1,7 @@ # Installation !!! important - vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0. + vLLM-Omni is released against the matching upstream vLLM major/minor version. The 0.29 development line requires vLLM 0.29.x; the source-install instructions pin vLLM 0.29.0. Published vLLM-Omni 0.28.0 wheels and images still require vLLM 0.28.x. vLLM-Omni supports the following hardware platforms: diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 7b8778296..e714a52bf 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,6 +34,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels +The source-install instructions target vLLM 0.29.0. Until vLLM-Omni 0.29.0 artifacts are published, the CUDA/ROCm pre-built examples below remain on the matching 0.28.0 release pair. + Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 5e6cd47c9..f6a1004a6 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -5,7 +5,7 @@ # --8<-- [end:requirements] # --8<-- [start:set-up-using-python] -vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x. +vLLM-Omni depends on the matching major/minor release of vLLM. The 0.29 development line uses vLLM 0.29.x. Published 0.28.0 wheels and images use vLLM 0.28.x. !!! note PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See for more details. @@ -18,6 +18,8 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** #### Installation of vLLM +These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. + vLLM-Omni is built based on vLLM. Please install it with command below. ```bash uv pip install vllm==0.28.0 --torch-backend=auto @@ -26,12 +28,12 @@ uv pip install vllm==0.28.0 --torch-backend=auto #### Installation of vLLM-Omni ```bash -uv pip install vllm-omni +uv pip install vllm-omni==0.28.0 ``` To run Gradio demos, also install the optional extras: ```bash -uv pip install 'vllm-omni[demo]' +uv pip install 'vllm-omni[demo]==0.28.0' ``` # --8<-- [end:pre-built-wheels] @@ -39,13 +41,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash -uv pip install vllm==0.28.0 --torch-backend=auto +uv pip install vllm==0.29.0 --torch-backend=auto ``` -The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.29.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +68,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.28.0 +git checkout v0.29.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.28.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.29.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.29.0/vllm-0.29.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -131,7 +133,7 @@ If you want to specify the base vLLM version: ```bash DOCKER_BUILDKIT=1 docker build \ -f docker/Dockerfile.cuda \ - --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \ + --build-arg BASE_IMAGE=vllm/vllm-openai:v0.29.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 e0f1614ae..d15bba2eb 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -11,6 +11,8 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. #### Installation of vLLM +These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. + vLLM-Omni is built based on vLLM. Please install it with command below. ```bash uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 @@ -22,7 +24,7 @@ uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc # we need to add --no-build-isolation as the torch # is not obtained from pypi, we have to install using the # torch installed in our environment -uv pip install vllm-omni +uv pip install vllm-omni==0.28.0 # Optional if want to run Qwen3 TTS uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm @@ -34,13 +36,13 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash -uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 +uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723 ``` -The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.29.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +60,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.28.0 +git checkout v0.29.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 036c2da17..8363e05ff 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.28.0 --torch-backend=auto +uv pip install vllm==0.29.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 +uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.29.0` and vLLM Omni `0.29.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py index d38380e13..4a8ea17e8 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -214,11 +214,20 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch): assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1 -def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size): +def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size, dtype=torch.float16): # Source of truth for the paged KV cache layout the backend must unpack. - from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + # vLLM 0.29 removed AttentionBackend.get_kv_cache_shape; the per-layer shape + # now comes from the spec, with the content dim expressed in bytes. + from vllm.v1.kv_cache_interface import FullAttentionSpec, compute_layer_kv_cache_shape_bytes - return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + ) + num_pages, num_heads, num_states, content_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks) + return (num_pages, num_heads, num_states, content_bytes // dtype.itemsize) def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 9011a1aef..7492a4b82 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -12,6 +12,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com from skip_ci import resolve_ci_decision # noqa: E402 from upload_pipeline import ( # noqa: E402 + CUDA_HF_TOKEN_EXPORT, _expand_mirror_hardwares, _get_mirror_hw_selector, _load_bootstrap_steps, @@ -143,6 +144,10 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0] assert container["image"].endswith("$BUILDKITE_COMMIT") assert container["resources"]["limits"]["nvidia.com/gpu"] == 1 + 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"] def test_mirror_hardwares_l4_preserves_explicit_retry() -> None: @@ -182,6 +187,33 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None: assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [ {"name": "swr-secret"}, ] + assert step["commands"] == ["pytest -sv tests/example"] + + +def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None: + for name in ( + "l4_1", + "l4_2", + "l4_3", + "l4_4", + "h100_1", + "h100_2", + "h100_3", + "h100_4", + "b200_1", + "b200_2", + "b200_3", + "b200_4", + ): + step = _expand_mirror_hardwares( + {"label": name, "mirror_hardwares": name, "commands": ["pytest -sv tests/example"]}, + ) + assert step is not None + container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0] + 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"][0] == CUDA_HF_TOKEN_EXPORT def _gpu_limit(step: dict) -> int: diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3e499dc38..40229f521 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -215,6 +215,16 @@ def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_s assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena" +@pytest.mark.parametrize("disabled", [True, False]) +def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled}) + assert config.stage_configs + engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model") + assert "disable_log_stats" not in engine_args + + def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime_bridge(): omni_config = _from_pipeline_key( "qwen3_tts", diff --git a/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py b/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py new file mode 100644 index 000000000..fb8546e2d --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The omni scheduler must stay on the v1 model runner, like every omni worker. + +vLLM 0.28 resolved this config to v1; 0.29 ends its fallback chain with v2. When +only the workers forced v1, the scheduler took its v2 fast path -- which omits +resumed-request token ids -- while the v1 runner still read them, so resuming a +request raised ``KeyError: `` in ``_update_states`` and killed the engine +core (Multi-Replica Startup Test). + +The default is upstream's to change, so this pins Omni's side of the contract +rather than the default itself. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Scheduler(OmniSchedulerMixin): + """Minimal carrier for the shared init hook both omni schedulers call.""" + + def __init__(self, *, use_v2_model_runner: bool) -> None: + self.use_v2_model_runner = use_v2_model_runner + self.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=False, + stage_id=0, + pooling_output_decoder=None, + ) + ) + self._init_omni_io_scheduling_state() + + +@pytest.mark.parametrize("initial", [True, False]) +def test_scheduler_ends_on_v1_model_runner(initial: bool) -> None: + """Whatever upstream resolves, the omni scheduler runs v1.""" + assert _Scheduler(use_v2_model_runner=initial).use_v2_model_runner is False + + +def test_v2_default_does_not_reach_the_resumed_request_fast_path() -> None: + """The v2 fast path is what dropped resumed-request token ids. + + ``OmniGenerationScheduler`` branches on ``use_v2_model_runner`` when it + assembles SchedulerOutput; forcing it false keeps the v1 assembly that + carries ``all_token_ids`` for resumed requests, which the v1 runner reads. + """ + scheduler = _Scheduler(use_v2_model_runner=True) + + assert not scheduler.use_v2_model_runner, ( + "scheduler would take the v2 fast path while omni workers run v1; " + "resumed requests then KeyError in _update_states" + ) diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py index f2002f1c6..aa8e07307 100644 --- a/tests/diffusion/diffusion_kv/test_block_tables.py +++ b/tests/diffusion/diffusion_kv/test_block_tables.py @@ -279,17 +279,17 @@ def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple] group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None)) return [[group], [group]], "cg-support", [4, 8] + # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the + # layout now rides on the resolved CacheConfig.kv_cache_layout instead. def init_cache( runner_kv_caches, forward_context, kv_cache_config, - attn_groups, device, - cache_dtype, kernel_block_sizes, vllm_config, ): - events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes)) + events.append(("cache", kv_cache_config, kernel_block_sizes)) runner_kv_caches.extend(["cache-0", "cache-1"]) return {"layer-0": "tensor-0", "layer-1": "tensor-1"} diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py index 861f9564d..b5f914b11 100644 --- a/tests/diffusion/diffusion_kv/test_initialization.py +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None ) assert vllm_config.model_config.original_max_model_len == -1 - # Sixteen physical pages provide ``16 * block_size`` tokens. The native - # backend chooses the block geometry, so keep this assertion backend + # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29 + # subtracts the null block that BlockPool permanently holds back before + # sizing (check_enough_kv_cache_memory), so only 15 pages are usable. The + # native backend chooses the block geometry, so keep this assertion backend # independent instead of assuming the generic 16-token block size. - expected_max_model_len = min(model_limit, spec.block_size * 16) + expected_max_model_len = min(model_limit, spec.block_size * (16 - 1)) assert vllm_config.model_config.max_model_len == expected_max_model_len diff --git a/tests/diffusion/diffusion_kv/test_layout.py b/tests/diffusion/diffusion_kv/test_layout.py new file mode 100644 index 000000000..713849e95 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_layout.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Physical KV cache layout selection and propagation for the Diffusion cache.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from vllm.config import CacheConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + DEFAULT_DIFFUSION_KV_CACHE_LAYOUT, + adopt_kv_cache_layout, + assert_backend_layout_supported, + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _config(): + return SimpleNamespace(cache_config=CacheConfig()) + + +class _DenseBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return False + + +class _BlockStrideBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + +def test_default_layout_is_not_block_outermost() -> None: + """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False.""" + vllm_config = _config() + layout = resolve_diffusion_kv_cache_layout(vllm_config) + + assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + assert layout.is_block_outermost is False + assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout + + +def test_block_stride_backend_selects_block_outermost_layout() -> None: + layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True) + + assert layout.is_block_outermost is True + + +def test_resolution_is_idempotent() -> None: + vllm_config = _config() + + assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config) + + +def test_contradicting_requirement_raises_instead_of_misaddressing() -> None: + vllm_config = _config() + resolve_diffusion_kv_cache_layout(vllm_config) + + with pytest.raises(ValueError, match="contradicts the attention"): + resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + +def test_preset_layout_is_honoured() -> None: + vllm_config = _config() + vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name + + layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + assert layout is KVCacheLayout.BLNHC + + +def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None: + """Workers collect specs before the control plane hands them the layout.""" + vllm_config = _config() + + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) # unresolved: nothing to check + + resolve_diffusion_kv_cache_layout(vllm_config) + assert_backend_layout_supported(vllm_config, _DenseBackend) + assert_backend_layout_supported(vllm_config, None) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) + + +def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None: + """Diffusion workers receive only a KVCacheConfig, so the name rides on it.""" + worker = _config() + shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name) + + assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC + assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC + + +def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None: + worker = _config() + + assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + + +def test_worker_disagreeing_with_the_control_plane_raises() -> None: + worker = _config() + resolve_diffusion_kv_cache_layout(worker) + + with pytest.raises(ValueError, match="disagrees with the layout"): + adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)) + + +def test_block_stride_requirement_propagates_control_plane_to_worker() -> None: + """A block-stride backend must reach the worker as a block-outermost layout. + + Covers the whole seam: the control plane resolves with the requirement, the + name rides on each rank-local KVCacheConfig, the worker adopts it, and the + backend contract check then passes instead of raising. + """ + control_plane = _config() + layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True) + assert layout.is_block_outermost is True + + shipped = SimpleNamespace(kv_cache_layout=layout.name) + + worker = _config() + assert adopt_kv_cache_layout(worker, shipped) is layout + # The worker's own backend requirement now agrees with what it was handed. + assert_backend_layout_supported(worker, _BlockStrideBackend) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(worker, _DenseBackend) + + +def test_tensor_strides_come_from_the_upstream_derivation() -> None: + spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) + + tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"]) + + assert tensor.layers == ["layer0", "layer1"] + assert tensor.size == spec.page_size_bytes * 2 * 4 + # Layer-outermost: each layer owns a contiguous run of its own blocks. + assert tensor.layer_stride == spec.page_size_bytes * 4 + assert tensor.block_stride == spec.page_size_bytes diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py index 5152c2a67..ed39f6523 100644 --- a/tests/diffusion/diffusion_kv/test_manager.py +++ b/tests/diffusion/diffusion_kv/test_manager.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest @@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig: ) return KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 1f242420d..535c1b411 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -15,6 +15,7 @@ import torch from torch import nn from vllm.v1.kv_cache_interface import FullAttentionSpec +from tests.helpers.kv_layout import layout_for_backend from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module @@ -1239,7 +1240,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads assert native_layer.num_heads == 4 assert native_layer.num_kv_heads == 2 assert native_layer.spec.num_kv_heads == 2 - assert native_layer.spec.indexes_kv_by_block_stride is True + # 0.29: the flag no longer rides on the spec, it selects the layout. + assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True assert native_layer._q_scale_float == 1.0 assert native_layer._k_scale_float == 1.0 assert native_layer._v_scale_float == 1.0 diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py index ef5d75f4d..d890dcc7b 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -10,12 +10,14 @@ import torch from torch import nn from torch.nn import functional as F from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache from vllm.v1.worker.gpu.block_table import BlockTables +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, DiffusionPagedAttentionLayerAdapter, @@ -88,12 +90,7 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N canonical_spec = native_layer.spec kv_cache_config = KVCacheConfig( num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor( - size=canonical_spec.page_size_bytes * 4, - shared_by=[_LAYER_NAME], - ) - ], + kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])], kv_cache_groups=[ KVCacheGroupSpec( layer_names=[_LAYER_NAME], @@ -101,6 +98,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N ) ], ) + # 0.29 reads the resolved physical layout when allocating; the diffusion + # control plane pins it, so this standalone harness must too. + resolve_diffusion_kv_cache_layout(vllm_config) with set_current_vllm_config(vllm_config): attn_groups, _, kernel_block_sizes = init_attn_backend( kv_cache_config, @@ -108,13 +108,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N device, ) runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now + # comes from the resolved CacheConfig.kv_cache_layout inside the call. init_kv_cache( runner_kv_caches, vllm_config.compilation_config.static_forward_context, kv_cache_config, - attn_groups, device, - vllm_config.cache_config.cache_dtype, kernel_block_sizes, vllm_config, ) diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py index ae5609b2a..ac9185e09 100644 --- a/tests/diffusion/diffusion_kv/test_request.py +++ b/tests/diffusion/diffusion_kv/test_request.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -10,10 +10,10 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager: group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec) config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])], kv_cache_groups=[group], ) return KVCacheManager( diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py index 9485ed046..1ccef7a9f 100644 --- a/tests/diffusion/diffusion_kv/test_spec_discovery.py +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -8,9 +8,10 @@ from unittest.mock import Mock import pytest import torch from torch import nn -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module +from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner @@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None: assert spec.num_kv_heads == 2 assert spec.head_size == 8 assert spec.dtype is torch.bfloat16 - assert spec.indexes_kv_by_block_stride is True + # vLLM 0.29 moved the block-stride decision off the spec onto the resolved + # physical layout; the backend's declaration must still select one. + assert layout_for_backend(_Backend).is_block_outermost is True assert spec.non_causal is True @@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)], ) @@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/layers/test_adalayernorm.py b/tests/diffusion/layers/test_adalayernorm.py index 5e41b7a26..94cadf5ac 100644 --- a/tests/diffusion/layers/test_adalayernorm.py +++ b/tests/diffusion/layers/test_adalayernorm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for shared AdaLayerNorm layers used by FLUX and other models.""" import os @@ -46,7 +46,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py index 284049347..d7d26350d 100644 --- a/tests/diffusion/layers/test_rope_broadcast.py +++ b/tests/diffusion/layers/test_rope_broadcast.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Shape broadcasting regression tests for RotaryEmbedding.forward_native. @@ -16,7 +16,11 @@ import types import pytest import torch -from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd +from vllm_omni.diffusion.layers.rope import ( + RotaryEmbedding, + RotaryEmbeddingWan, + apply_rotary_emb_mindiesd, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -239,6 +243,25 @@ class TestRotaryEmbeddingNativeShapeRegression: assert not torch.isnan(output).any() +def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None: + """Source/CPU installs need not contain the wheel-vendored CUDA kernel.""" + x = torch.randn(1, 5, 2, 8) + + rope = RotaryEmbedding(is_neox_style=True) + rope.apply_rotary_emb_vllm_flash_attn = None + output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4)) + assert output.shape == x.shape + + wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + wan_rope.apply_rotary_emb_vllm_flash_attn = None + output = wan_rope.forward_cuda( + x, + torch.randn(1, 5, 1, 4), + torch.randn(1, 5, 1, 4), + ) + assert output.shape == x.shape + + def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None: """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout.""" captured: dict[str, object] = {} diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 4296c1cde..43e2f12f4 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -59,7 +59,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/dreamzero/test_qkv_fusion.py b/tests/diffusion/models/dreamzero/test_qkv_fusion.py index 16b78693f..d67a01e48 100644 --- a/tests/diffusion/models/dreamzero/test_qkv_fusion.py +++ b/tests/diffusion/models/dreamzero/test_qkv_fusion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Equivalence test for DreamZero self-attn QKV fusion. @@ -52,7 +52,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py index 8ab046f0c..1a6170ca2 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_quantization.py +++ b/tests/diffusion/models/glm_image/test_glm_image_quantization.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for GLM-Image quantization support (W4A16/AutoRound). These tests verify that the GLM-Image DiT transformer correctly accepts and uses @@ -69,7 +69,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py index 20024abe7..62bcf9ce8 100644 --- a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py +++ b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace import pytest @@ -64,6 +64,13 @@ def setup_tp_group(monkeypatch, mocker): lambda: 0, ) + # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever + # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path. + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda", + lambda: False, + ) + mock_tp_group = mocker.MagicMock() mock_tp_group.world_size = 2 mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group) @@ -94,7 +101,7 @@ def setup_tp_group(monkeypatch, mocker): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) with set_current_vllm_config(VllmConfig(device_config=device_config)): diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py index 52f1b6aa3..4f3a81491 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py @@ -1,12 +1,39 @@ import pytest import torch -from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + QwenEmbedRope, + _apply_qwen_image_rotary_emb, +) from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("noncontiguous", [False, True]) +def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous): + from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen + + generator = torch.Generator().manual_seed(42) + x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype) + if noncontiguous: + x = x.transpose(1, 2).contiguous().transpose(1, 2) + angles = torch.randn(17, 64, generator=generator) + freqs = torch.polar(torch.ones_like(angles), angles) + expected = apply_rotary_emb_qwen(x, freqs, use_real=False) + actual = _apply_qwen_image_rotary_emb(x, freqs) + assert actual.dtype == dtype + assert actual.shape == x.shape + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + if dtype == torch.bfloat16: + # A regression to early BF16 frequency rounding must be observable. + rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float()) + rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs) + assert not torch.equal(actual, rounded) + + def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count(): prompt_embeds = torch.zeros(2, 64, 8) prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py index 0bc2ddb8e..9e6404b21 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os @@ -22,7 +22,7 @@ def _init_distributed(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29501") diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py index 3f4c62075..6a0f0548e 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py @@ -41,8 +41,8 @@ def tp1_group(): os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29501") - init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") - initialize_model_parallel() + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo") + initialize_model_parallel(backend="gloo") yield cleanup_dist_env_and_memory() @@ -54,7 +54,7 @@ def force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) @@ -482,8 +482,12 @@ class _LockstepP2PDist: def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None: """Two lockstep rank threads over uneven frame shards must reproduce the dense tiny-transformer output, and must actually have sharded the work.""" + from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + # These CPU tensor tests must not probe a CUDA device during construction. + mocker.patch("vllm_omni.diffusion.attention.selector._cached_get_backend_cls", return_value=SDPABackend) + torch.manual_seed(3) model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval() for _, param in sorted(model.named_parameters()): diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py index 41ca717d6..a44c90146 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and sharded GLUMB temporal conv. @@ -51,7 +51,7 @@ def force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py index b36acc68a..cca73fb08 100644 --- a/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py +++ b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch from transformers import T5Config @@ -54,6 +54,13 @@ def setup_tp_group(monkeypatch, mocker): lambda: 0, ) + # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever + # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path. + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda", + lambda: False, + ) + # TP group mock_tp_group = mocker.MagicMock() mock_tp_group.world_size = 2 diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py index 943411779..6aebd6b60 100644 --- a/tests/diffusion/quantization/test_component_routing.py +++ b/tests/diffusion/quantization/test_component_routing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for component routing for quantization.""" from unittest.mock import MagicMock @@ -372,7 +372,8 @@ class TestComponentApplyVllmMapper: "thinker.model": "thinker.language_model.model", "talker.model": "talker.language_model.model", } - ).get_unstacked_mapper() + # vLLM 0.29 renamed get_unstacked_mapper to get_rename_mapper. + ).get_rename_mapper() config.apply_vllm_mapper(mapper) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index d6d17c18f..20a2f0384 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import queue @@ -10,8 +10,9 @@ import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from pytest_mock import MockerFixture -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -104,7 +105,7 @@ def _initialize_paged_scheduler( ) config = KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) scheduler.initialize( diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index be3d7f84c..b12e9c1e7 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -164,6 +164,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FLOW_SHIFT) pipe.to("cuda") + _set_matched_attention_backend(pipe) _diffusers_dummy_run(pipe) diff --git a/tests/e2e/online_serving/test_dots_tts_expansion.py b/tests/e2e/online_serving/test_dots_tts_expansion.py index 34bc54b8e..e5df55a33 100644 --- a/tests/e2e/online_serving/test_dots_tts_expansion.py +++ b/tests/e2e/online_serving/test_dots_tts_expansion.py @@ -85,5 +85,7 @@ def test_text_only_nonstreaming_pcm_001(omni_server, openai_client) -> None: "voice": "default", "min_audio_bytes": _MIN_AUDIO_BYTES, "min_hnr_db": -2.0, + # dots.tts-soar emits native 48 kHz PCM, which has no rate header. + "expected_sample_rate": 48_000, } openai_client.send_audio_speech_request(request_config) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index e6f7f8a8b..27677bd54 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1593,73 +1593,58 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path): assert "model_path_resolver" not in engine_args -def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): +def test_build_stage0_input_processor_uses_omni_renderer_subclass(monkeypatch): + from vllm.renderers import BaseRenderer + import vllm_omni.engine.stage_init_utils as init_mod + from vllm_omni.inputs.preprocess import OmniRenderer, omni_renderer_cls - class DummyInputProcessor: - def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config - self.renderer = renderer or object() - self.input_preprocessor = None + class _Base(BaseRenderer): + def __init__(self, config, tokenizer): + self.config, self.tokenizer = config, tokenizer - class DummyOmniInputPreprocessor: - def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config - self.renderer = renderer + def render_messages(self, messages, params): # pragma: no cover - abstract stub + raise NotImplementedError - monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) - - input_processor = build_stage0_input_processor( - types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) - ) - - assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) - assert input_processor.input_preprocessor.renderer is input_processor.renderer - - -def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped( - monkeypatch, -): - import vllm_omni.engine.stage_init_utils as init_mod - - token_only_renderer = object() + config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) + built = omni_renderer_cls(_Base)(config, "tok") seen = {} class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - if renderer is None: - raise AssertionError("skip_tokenizer_init must supply a renderer") seen["renderer"] = renderer self.renderer = renderer - self.input_preprocessor = None - - class DummyOmniInputPreprocessor: - def __init__(self, vllm_config, renderer=None): - seen["preprocessor_renderer"] = renderer monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr( - init_mod, - "_build_token_only_renderer", - lambda _config: token_only_renderer, - ) - monkeypatch.setattr( - init_mod, - "OmniInputPreprocessor", - DummyOmniInputPreprocessor, - ) + monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None) + processor = build_stage0_input_processor(config) + assert seen["renderer"] is built + assert isinstance(processor.renderer, OmniRenderer) + assert isinstance(processor.renderer, _Base) + assert not hasattr(processor, "input_preprocessor") + +def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(monkeypatch): + import vllm_omni.engine.stage_init_utils as init_mod + + original = object() + + class DummyInputProcessor: + def __init__(self, vllm_config, renderer=None): + assert renderer is original + self.renderer = renderer + + def _must_not_resolve(_cfg): + raise AssertionError("tokenizer must not be resolved when skip_tokenizer_init=True") + + monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) + monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original) + monkeypatch.setattr(init_mod, "build_omni_renderer", _must_not_resolve) config = types.SimpleNamespace( - model_config=types.SimpleNamespace( - skip_tokenizer_init=True, - try_get_generation_config=lambda: {}, - ) + model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {}) ) - build_stage0_input_processor(config) - - assert seen["renderer"] is token_only_renderer - assert seen["preprocessor_renderer"] is token_only_renderer + processor = build_stage0_input_processor(config) + assert processor.renderer is original def test_inject_kv_stage_info_infers_sender_tp_topology(): diff --git a/tests/entrypoints/openai/test_batch_chat_completions.py b/tests/entrypoints/openai/test_batch_chat_completions.py index 9fc72b2c1..c3866dbc8 100644 --- a/tests/entrypoints/openai/test_batch_chat_completions.py +++ b/tests/entrypoints/openai/test_batch_chat_completions.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio from unittest.mock import AsyncMock @@ -12,7 +15,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionResponseChoice, ChatMessage, ) -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo +from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index a9bf76319..9c7eb9d89 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -248,7 +248,7 @@ class TestResolveAudioFormat: assert result == "pcm" def test_invalid_format_returns_error(self, serving_chat): - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse request = self._make_request({"format": "aac", "voice": "alloy"}) result = serving_chat._resolve_audio_format(request) @@ -256,7 +256,7 @@ class TestResolveAudioFormat: assert "aac" in result.error.message def test_all_supported_formats_accepted(self, serving_chat): - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse for fmt in SUPPORTED_CHAT_AUDIO_FORMATS: request = self._make_request({"format": fmt, "voice": "alloy"}) diff --git a/tests/entrypoints/openai_api/test_modality_validation.py b/tests/entrypoints/openai_api/test_modality_validation.py index fde0c0ca1..91198b2e2 100644 --- a/tests/entrypoints/openai_api/test_modality_validation.py +++ b/tests/entrypoints/openai_api/test_modality_validation.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for output modality validation in chat completions. Covers: @@ -13,7 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse pytestmark = [pytest.mark.core_model, pytest.mark.cpu] diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index 2fdf9f646..6e796f480 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" from types import SimpleNamespace @@ -19,7 +20,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_omni_chat_completion_response_metrics(): """Test OmniChatCompletionResponse metrics field works correctly.""" - from vllm.entrypoints.openai.engine.protocol import UsageInfo + from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm_omni.entrypoints.openai.protocol.chat_completion import ( OmniChatCompletionResponse, diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6a5ebcf7d..3922ab4a7 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -25,7 +25,7 @@ from fastapi.responses import JSONResponse from fastapi.testclient import TestClient from pydantic import ValidationError from pytest_mock import MockerFixture -from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py index e382b5e32..31717d0b7 100644 --- a/tests/entrypoints/openai_api/test_video_stream_handler.py +++ b/tests/entrypoints/openai_api/test_video_stream_handler.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Runtime behavior tests for VideoStreamHandler. Every test creates a mock WebSocket, drives handle_session through a @@ -281,7 +281,7 @@ class TestHandleQueryErrors: @pytest.mark.asyncio async def test_query_returns_error_response(self): """create_chat_completion returns ErrorResponse → error forwarded.""" - from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse err = ErrorResponse( error=ErrorInfo( diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 97ed2370b..87aeb787e 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for multi-modal streaming finish_reason behavior (commit 44c799bc). Verifies that the /v1/chat/completions streaming endpoint emits exactly one @@ -411,7 +412,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): @pytest.mark.asyncio async def test_audio_choice_error_response_is_not_iterated_as_choices(): """Any ErrorResponse from the audio path is skipped, not iterated.""" - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse serving_chat = build_serving_chat() request = make_request(modalities=["text", "audio"]) diff --git a/tests/helpers/kv_layout.py b/tests/helpers/kv_layout.py new file mode 100644 index 000000000..8351de8da --- /dev/null +++ b/tests/helpers/kv_layout.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared access to the production ``KVCacheTensor`` layout derivation. + +Tests deliberately go through the same code production uses +(``vllm_omni.diffusion.diffusion_kv.layout``) rather than recomputing strides, +so a regression in that derivation fails the suite instead of being masked by a +parallel test-only copy. +""" + +from types import SimpleNamespace + +from vllm.config import CacheConfig +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +__all__ = ["build_kv_cache_tensor", "layout_for_backend"] + + +def layout_for_backend(attn_backend) -> KVCacheLayout: + """The physical layout ``attn_backend``'s block-stride declaration selects. + + Before 0.29 a backend's ``indexes_kv_by_block_stride()`` was copied onto + ``AttentionSpec``; asserting on the resolved layout is the surviving way to + check that declaration still reaches the cache. + """ + + vllm_config = SimpleNamespace(cache_config=CacheConfig()) + return resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), + ) diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py index f1f2710b5..45add4b63 100644 --- a/tests/helpers/serving_chat.py +++ b/tests/helpers/serving_chat.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared scaffolding for OmniOpenAIServingChat unit tests. Provides factories for a minimal serving-chat instance, chat requests, and @@ -8,11 +9,11 @@ OmniRequestOutput objects, plus SSE stream collection/parsing utilities. import json from unittest.mock import MagicMock +from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponseStreamChoice, ) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.outputs import CompletionOutput, RequestOutput diff --git a/tests/inputs/test_mm_processor.py b/tests/inputs/test_mm_processor.py new file mode 100644 index 000000000..813d7f436 --- /dev/null +++ b/tests/inputs/test_mm_processor.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from transformers import BatchFeature + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("counts", [{}, {"audio": 1, "image": 0}]) +@pytest.mark.parametrize("original_text", [None, "", "voice cloning text"]) +def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, original_text): + 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(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 + kwargs = {"sampling_rate": 16000} + if original_text is not None: + kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text + before = dict(kwargs) + result = processor._apply_hf_processor_main(items, kwargs) + 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"]}, + {"sampling_rate": 16000}, + {}, + ) + assert dict(result) == {"features": "processed", "embedding": "preserve"} + assert kwargs == before diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index db8817cf2..3fb48a7cc 100644 --- a/tests/inputs/test_preprocess.py +++ b/tests/inputs/test_preprocess.py @@ -1,35 +1,157 @@ # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OmniInputPreprocessor._process_text routing.""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise Omni routing through the active upstream rendering pipeline.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest -from pytest_mock import MockerFixture +import torch +from vllm.inputs import tokens_input +from vllm.pooling_params import PoolingParams +from vllm.renderers import BaseRenderer +from vllm.renderers.params import TokenizeParams +from vllm.v1.engine.input_processor import InputProcessor + +from vllm_omni.inputs import preprocess as preprocess_mod +from vllm_omni.inputs.preprocess import OmniRenderer, build_omni_renderer, omni_renderer_cls pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -class TestProcessTextMmProcessorKwargsRouting: - """Presence-based routing: an explicit empty ``mm_processor_kwargs`` dict - on the prompt still routes through ``_process_multimodal``. Required for - AR-based image generation (e.g. GLM-Image t2i) where the HF processor - supplies the image-generation scaffold from its own defaults and no extra - kwargs are needed from the caller. - """ +class _Renderer(BaseRenderer): + def __init__(self): + self.model_config = SimpleNamespace(is_encoder_decoder=False, enable_prompt_embeds=True) + self.tokenizer = None + self.default_cmpl_tok_params = TokenizeParams(max_total_tokens=128) + self._tokenize_prompt = lambda prompt, params: {**prompt, "prompt_token_ids": [1, 2, 3]} + self._tokenize_prompt_async = AsyncMock(side_effect=self._tokenize_prompt) + self._process_multimodal = Mock(return_value=tokens_input([1, 2, 3, 99])) + self._process_multimodal_async = AsyncMock(side_effect=self._process_multimodal) - @pytest.fixture - def preprocessor(self, mocker: MockerFixture): - from vllm_omni.inputs.preprocess import OmniInputPreprocessor + def render_messages(self, messages, params): + return messages, {"prompt": "hello"} - instance = object.__new__(OmniInputPreprocessor) - instance._process_multimodal = mocker.MagicMock(return_value={}) - instance._tokenize_prompt = mocker.MagicMock(return_value=[1, 2, 3]) - return instance - def test_empty_mm_processor_kwargs_routes_to_multimodal(self, preprocessor): - preprocessor._process_text({"prompt": "hello", "mm_processor_kwargs": {}}) - assert preprocessor._process_multimodal.called - assert not preprocessor._tokenize_prompt.called +@pytest.fixture +def renderer(): + return omni_renderer_cls(_Renderer)() - def test_missing_mm_processor_kwargs_routes_to_tokenize(self, preprocessor): - preprocessor._process_text({"prompt": "hello"}) - assert preprocessor._tokenize_prompt.called - assert not preprocessor._process_multimodal.called + +@pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}]) +@pytest.mark.parametrize("tokenized", [False, True]) +def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, tokenized): + # Keep real process_inputs/render_cmpl; stub unrelated config validation. + processor = object.__new__(InputProcessor) + processor.renderer = renderer + processor.model_config = renderer.model_config + processor.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_size_local=1, + local_engines_only=False, + ) + ) + processor._validate_params = Mock() + processor._validate_lora = Mock() + processor._validate_model_inputs = Mock() + prompt: dict[str, object] = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"} + prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt") + request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",)) + assert request.prompt_token_ids == [1, 2, 3, 99] + assert request.cache_salt == "salt" + renderer._process_multimodal.assert_called_once_with( + [1, 2, 3], + {}, + mm_processor_kwargs=kwargs, + mm_uuids=None, + skip_mm_cache=False, + ) + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"]) +def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind): + prompt: dict[str, object] = { + "prompt": "hello", + "additional_information": {"speaker": 1}, + "model_intermediate_buffer": {"x": 2}, + } + if kind == "tokens": + prompt["prompt_token_ids"] = [1, 2, 3] + elif kind == "media": + prompt["multi_modal_data"] = {"image": "image"} + elif kind == "kwargs": + prompt["mm_processor_kwargs"] = {} + elif kind == "embeds": + prompt["prompt_embeds"] = torch.ones(3, 4) + if async_mode: + (result,) = asyncio.run(renderer.render_cmpl_async([prompt], skip_mm_cache=True)) + else: + (result,) = renderer.render_cmpl([prompt], skip_mm_cache=True) + assert result["additional_information"] == {"speaker": 1} + assert result["model_intermediate_buffer"] == {"x": 2} + assert result["prompt"] == "hello" + if kind in ("media", "kwargs"): + assert result["prompt_token_ids"] == [1, 2, 3, 99] + assert renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True + else: + renderer._process_multimodal.assert_not_called() + if kind == "embeds": + torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"]) + else: + assert result["prompt_token_ids"] == [1, 2, 3] + + +def test_omni_renderer_is_a_real_subclass_created_once(): + cls = omni_renderer_cls(_Renderer) + assert issubclass(cls, OmniRenderer) and issubclass(cls, _Renderer) + assert omni_renderer_cls(_Renderer) is cls + assert omni_renderer_cls(cls) is cls + assert cls.__name__ == "Omni_Renderer" + instance = cls() + assert isinstance(instance, _Renderer) + assert not hasattr(instance, "_renderer") + + +def test_concrete_renderer_overrides_are_not_suppressed(): + """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides.""" + + class _Custom(_Renderer): + def render_cmpl(self, prompts, *args, **kwargs): + return ["custom"] + + assert omni_renderer_cls(_Custom)().render_cmpl([{"prompt": "x"}]) == ["custom"] + + +class _CtorRenderer(BaseRenderer): + def __init__(self, config, tokenizer): + self.config = config + self.tokenizer = tokenizer + + def render_messages(self, messages, params): + raise NotImplementedError + + +def test_build_omni_renderer_with_explicit_class_skips_registry(monkeypatch): + monkeypatch.setattr(preprocess_mod, "cached_tokenizer_from_config", Mock(side_effect=AssertionError("resolved"))) + config = SimpleNamespace(model_config=SimpleNamespace()) + built = build_omni_renderer(config, renderer_cls=_CtorRenderer, tokenizer=None) + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.config is config and built.tokenizer is None + + +def test_build_omni_renderer_resolves_like_upstream(monkeypatch): + config = SimpleNamespace(model_config=SimpleNamespace()) + tokenizer = object() + monkeypatch.setattr( + preprocess_mod, "cached_tokenizer_from_config", lambda mc: tokenizer if mc is config.model_config else None + ) + monkeypatch.setattr(preprocess_mod, "tokenizer_args_from_config", lambda mc: ("fake-mode", None)) + load_cls = Mock(return_value=_CtorRenderer) + monkeypatch.setattr(preprocess_mod.RENDERER_REGISTRY, "load_renderer_cls", load_cls) + built = build_omni_renderer(config) + load_cls.assert_called_once_with("fake-mode") + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.tokenizer is tokenizer diff --git a/tests/model_executor/models/glm_image/test_glm_image_ar.py b/tests/model_executor/models/glm_image/test_glm_image_ar.py index 5ccc9fc62..ae1060839 100644 --- a/tests/model_executor/models/glm_image/test_glm_image_ar.py +++ b/tests/model_executor/models/glm_image/test_glm_image_ar.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for GLM-Image AR model: DataParser, processor, and M-RoPE.""" import importlib.util @@ -91,6 +91,41 @@ GlmImageRotaryEmbedding = _ar_mod.GlmImageRotaryEmbedding pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("num_images", [0, 1, 2]) +@pytest.mark.parametrize("preformatted", [False, True]) +def test_generation_prompt_ids_preserve_hf_target_scaffold(num_images, preformatted): + from types import SimpleNamespace + + from vllm.multimodal.parse import MultiModalDataItems + from vllm.multimodal.processing import BaseMultiModalProcessor, ProcessorInputs + + processor = object.__new__(GlmImageMultiModalProcessor) + # Only the public token attributes of GlmImageProcessor are needed; the + # expected suffix below is what HF's own formatter produces for 512x768. + hf_processor = SimpleNamespace(grid_bos_token="", grid_eos_token="", bos_token="") + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) + processor.info = SimpleNamespace( + get_tokenizer=lambda: tokenizer, + get_hf_processor=lambda: hf_processor, + get_hf_config=lambda: _make_hf_config(), + ) + items = MagicMock(spec=MultiModalDataItems) + items.get_all_counts.return_value = {"image": num_images} + suffix = "16 24" + ("13 19" if not num_images else "") + "" + text_ids = list(b"a red bird") + expected = [167855] * num_images + text_ids + list(suffix.encode()) + inputs = ProcessorInputs( + prompt=expected if preformatted else text_ids, + mm_data_items=items, + hf_processor_mm_kwargs={"target_h": 512, "target_w": 768}, + ) + with patch.object(BaseMultiModalProcessor, "apply", side_effect=lambda inputs, timing: inputs) as upstream: + result = processor.apply(inputs, MagicMock()) + assert result.prompt == expected + assert inputs.prompt == (expected if preformatted else text_ids) + upstream.assert_called_once() + + def test_glm_image_ar_omits_hidden_pooler_payload(): assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py index d25b58e48..48d3bd9a3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path. Two defects made the OpenAI ``audio_embeds`` content part unusable with this model: @@ -71,6 +71,12 @@ class _IdentityTokenizer: def decode(self, tokens: str) -> str: return tokens + def convert_tokens_to_ids(self, token: str) -> int: + # vLLM 0.29 prompt updates select placeholder positions by token id, + # so the processor resolves ```` once; the value is irrelevant here. + assert token == "" + return 0 + def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo: processor = _FakeProcessor() diff --git a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py index 3e8ee2a63..eaa2ff2b3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py +++ b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for MiniCPM-o 4.5 >30s audio preprocessing. Covers: @@ -43,6 +43,34 @@ _FEAT_DIM = 128 _CHUNK_FRAMES = 3000 +@pytest.mark.parametrize("version", [(2, 5), (4, 5)]) +def test_prompt_batching_delegates_upstream_without_tts_kwargs(version): + from unittest.mock import MagicMock + + info = SimpleNamespace( + get_model_version=lambda: version, + get_hf_processor=MagicMock(return_value=object()), + ctx=SimpleNamespace(call_hf_processor=MagicMock()), + ) + info.ctx.call_hf_processor.side_effect = ( + [{"features": [1]}, {"features": [2]}] if version == (2, 5) else [{"features": [1, 2]}] + ) + kwargs = {"use_tts": True, "sampling_rate": 16000} + result = MiniCPMO45OmniLLMMultiModalProcessor._call_hf_processor_on_prompts( + SimpleNamespace(info=info), ["one", "two"], {"audios": [10, 20]}, kwargs, out_keys={"features"} + ) + assert result == {"features": [1, 2]} + assert kwargs == {"use_tts": True, "sampling_rate": 16000} + for call in info.ctx.call_hf_processor.call_args_list: + assert call.args[2] == {"sampling_rate": 16000} + payloads = [call.args[1] for call in info.ctx.call_hf_processor.call_args_list] + assert payloads == ( + [{"text": "one", "audios": 10}, {"text": "two", "audios": 20}] + if version == (2, 5) + else [{"text": ["one", "two"], "audios": [10, 20]}] + ) + + class TestMiniCPMOFieldConfig: def test_short_audios_stay_batched(self) -> None: # two <=30s audios: one feature entry + one 1-element lens per audio @@ -99,10 +127,10 @@ def _make_processor( processor.info = SimpleNamespace(audio_pattern="()") processor.data_parser = MiniCPMOMultiModalDataParser(target_sr=_SAMPLE_RATE) - def fake_base_call_hf_processor(prompts, mm_data, mm_kwargs, tok_kwargs, *, out_keys): + def fake_call_hf_processor_on_prompts(prompts, mm_data, mm_kwargs, *, out_keys): return {key: fake_hf_outputs[key] for key in out_keys} - processor._base_call_hf_processor = fake_base_call_hf_processor + processor._call_hf_processor_on_prompts = fake_call_hf_processor_on_prompts return processor @@ -120,7 +148,7 @@ class TestProcessAudiosUnpadding: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}) features = result["audio_features"] assert [feat.shape[-1] for feat in features] == chunk_lens @@ -138,7 +166,7 @@ class TestProcessAudiosUnpadding: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}) widths = [feat.shape[-1] for feat in result["audio_features"]] assert widths == [_CHUNK_FRAMES, 1200, _CHUNK_FRAMES, 800] @@ -151,7 +179,7 @@ class TestProcessAudiosUnpadding: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}) features = result["audio_features"] assert len(features) == 1 diff --git a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py new file mode 100644 index 000000000..9c6533c76 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniVoice processor: the prompt is tokenized exactly once per request. + +Upstream ``_cached_apply_hf_processor`` calls ``_apply_hf_processor_main`` for +every request, including text-only requests and requests whose reference audio +is already cached (empty ``mm_items``). These tests pin that only ``apply`` +touches the text tokenizer and that the media path only encodes the audio. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ProcessorInputs + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.omnivoice.omnivoice import OmniVoiceMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_PROMPT_IDS = [5, 6, 7] + + +def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProcessor: + processor = object.__new__(OmniVoiceMultiModalProcessor) + tokenizer = SimpleNamespace(decode=Mock(return_value="hello world")) + processor.info = SimpleNamespace( + ctx=SimpleNamespace( + get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(sample_rate=24000), + model_config=SimpleNamespace(model="fake-model"), + ), + get_tokenizer=lambda: tokenizer, + ) + processor.text_tokenizer = Mock(return_value=SimpleNamespace(input_ids=torch.tensor([_PROMPT_IDS]))) + processor.audio_tokenizer = SimpleNamespace(encode=Mock(return_value=torch.zeros(8, 4, dtype=torch.long))) + processor.feature_extractor = None + monkeypatch.setattr(processor, "_ensure_cached_runtime_components", lambda *_args, **_kwargs: None) + return processor + + +def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace: + return SimpleNamespace( + get_all_counts=lambda: counts, + select=lambda _keys: mm_data, + ) + + +def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _make_processor(monkeypatch) + captured: dict[str, ProcessorInputs] = {} + + def fake_parent_apply(self, inputs, timing_ctx): + captured["inputs"] = inputs + return "sentinel" + + monkeypatch.setattr(OmniMultiModalProcessor, "apply", fake_parent_apply) + inputs = ProcessorInputs( + prompt=[1, 2], + mm_data_items=MultiModalDataItems({}), + hf_processor_mm_kwargs={"ref_text": "reference"}, + ) + + assert processor.apply(inputs, timing_ctx=None) == "sentinel" + assert captured["inputs"].prompt == _PROMPT_IDS + assert processor.text_tokenizer.call_count == 1 + (text_prompt,), _ = processor.text_tokenizer.call_args + assert text_prompt.endswith("<|text_start|>reference hello world<|text_end|>") + assert captured["inputs"].hf_processor_mm_kwargs["ref_text"] == "reference" + + +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"}) + + assert tuple(out["ref_audio_tokens"].shape) == (8, 4) + assert list(out["ref_audio_len"]) == [4] + assert "input_ids" not in out + assert processor.audio_tokenizer.encode.call_count == 1 + assert processor.text_tokenizer.call_count == 0 + + +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"}) + + assert len(out.keys()) == 0 + assert processor.text_tokenizer.call_count == 0 + assert processor.audio_tokenizer.encode.call_count == 0 + + +def test_legacy_call_hf_processor_is_gone() -> None: + assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor") 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 d88d918bc..54efc99e4 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 @@ -266,8 +266,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): inputs = MagicMock() inputs.mm_data_items = mm_data_items inputs.hf_processor_mm_kwargs = {"use_audio_in_video": True} - inputs.tokenization_kwargs = {} - inputs.prompt = [1, 2, 3] inputs.get_mm_hashes = MagicMock(return_value={"video": ["v0"], "audio": ["a0"]}) timing_ctx = MagicMock() @@ -307,7 +305,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): ) fake_self._paired_cache_keys = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys fake_self.info.parse_mm_data = MagicMock(return_value=SimpleNamespace(name="missing")) - fake_self._apply_hf_processor_main = MagicMock(return_value=([9], {"video_grid_thw": object()}, False)) + fake_self._apply_hf_processor_main = MagicMock(return_value={"video_grid_thw": object()}) 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 @@ -316,13 +314,12 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs", return_value={"video": [object()], "audio": [object()]}, ): - prompt_ids, mm_info, _ = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor( + mm_info = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor( fake_self, inputs, timing_ctx, ) - assert prompt_ids == [9] # Both modalities were reprocessed as one unit after pair-miss expansion. assert fake_self.info.parse_mm_data.call_args.args[0] == { "video": ["video-0"], @@ -623,10 +620,17 @@ def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): + from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as UpstreamQwen3Processor, + ) + assert issubclass( Qwen3OmniMoeThinkerMultiModalProcessor, Qwen2_5OmniThinkerMultiModalProcessor, ) + assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data + mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__ + assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor) def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py index 75d81eed5..4ff6d5096 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py @@ -113,7 +113,9 @@ def _save_peft_adapter(tmp_path: Path) -> Path: return adapter_dir -def test_peft_expert_lora_loads_into_thinker(tmp_path: Path): +def test_peft_expert_lora_loads_into_thinker(tmp_path: Path, monkeypatch): + # This CPU harness tests weight mapping, not CUDA pinned-memory allocation. + monkeypatch.setattr("vllm.lora.lora_model.PIN_MEMORY", False) adapter_dir = _save_peft_adapter(tmp_path) peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128) lora_model = LoRAModel.from_local_checkpoint( diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 45b40d026..64f9ac8de 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace from unittest.mock import patch @@ -157,7 +157,10 @@ def _load_weights_noop(model: Qwen3TTSCode2Wav) -> set[str]: def __init__(self, *_: object, **__: object): pass - def load_weights(self, _weights: object) -> set[str]: + def load_weights(self, _weights: object, *, mapper: object | None = None) -> set[str]: + # vLLM 0.29 dropped skip_prefixes from AutoWeightsLoader; the model + # now passes the stage exclusions as a WeightsMapper instead. + del mapper return {"decoder.fake_weight"} with ( diff --git a/tests/model_executor/models/test_weight_loader.py b/tests/model_executor/models/test_weight_loader.py new file mode 100644 index 000000000..34d1cd28a --- /dev/null +++ b/tests/model_executor/models/test_weight_loader.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from torch import nn +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TinyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.kept = nn.Linear(1, 1, bias=False) + self.skipped = nn.Linear(1, 1, bias=False) + self.rotary_embed = nn.Module() + self.rotary_embed.register_parameter("inv_freq", nn.Parameter(torch.zeros(1))) + + +def test_upstream_mapper_preserves_stage_exclusions_after_renaming() -> None: + model = _TinyModel() + model.skipped.weight.data.zero_() + + loaded = AutoWeightsLoader(model).load_weights( + [ + ("checkpoint.kept.weight", torch.ones(1, 1)), + ("checkpoint.skipped.weight", torch.full((1, 1), 2.0)), + ("rotary_embed.inv_freq", torch.ones(1)), + ], + mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""}) + | WeightsMapper(orig_to_new_prefix={"skipped.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None}), + ) + + assert loaded == {"kept.weight"} + torch.testing.assert_close(model.kept.weight, torch.ones(1, 1)) + torch.testing.assert_close(model.skipped.weight, torch.zeros(1, 1)) + torch.testing.assert_close(model.rotary_embed.inv_freq, torch.zeros(1)) diff --git a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py index 0e9f23de6..985c64114 100644 --- a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py +++ b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by ``vllm_omni.patch``. @@ -198,7 +198,7 @@ def test_reload_does_not_nest_wrapper(): print("NESTED" if nested else "FLAT") """ ) - result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600) if result.stdout.strip().startswith("SKIP:"): pytest.skip(result.stdout.strip()) assert "FLAT" in result.stdout, ( @@ -265,9 +265,12 @@ def test_env_var_escape_hatch_disables_install(env_value, acceptable): ) env = os.environ.copy() env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value - # 120s rather than 60s: cold CI containers can take 30-50s just to - # import vllm_omni → vllm → torch before the subprocess does any work. - result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env) + # Generous timeout: cold CI containers can take 30-50s just to import + # vllm_omni → vllm → torch, and heavily loaded shared machines have + # been observed taking >2min for the same import under disk I/O + # contention. The subprocess does a single import + flag print, so 600s + # only bounds pathological hangs, not normal runtime. + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600, env=env) if result.stdout.strip().startswith("SKIP:"): pytest.skip("modelopt not available in subprocess") # Exact token match on the final stdout line — not a substring test, since diff --git a/tests/model_executor/test_upstream_prompt_update_api.py b/tests/model_executor/test_upstream_prompt_update_api.py new file mode 100644 index 000000000..da31980c2 --- /dev/null +++ b/tests/model_executor/test_upstream_prompt_update_api.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Guard that Omni only calls prompt-update helpers upstream still provides. + +vLLM 0.29 removed ``PromptUpdateDetails.select_text`` in favour of the +token-oriented ``select_token_id`` / ``select_token_ids``. Three MiniCPM-o call +sites survived the rebase and only failed at engine-core startup, where the +traceback surfaced as an opaque "Engine core initialization failed" across eight +CI jobs. + +These helpers are referenced from replacement closures that no unit test drives, +so a static check is what actually catches the next removal. It resolves the +attribute against the installed class rather than a hard-coded list, so it +follows upstream automatically. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +from vllm.multimodal.processing.processor import PromptUpdateDetails + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_OMNI_ROOT = Path(__file__).resolve().parents[2] / "vllm_omni" +_GUARDED = "PromptUpdateDetails" + + +def _referenced_attributes() -> dict[str, set[str]]: + """Map ``PromptUpdateDetails.`` -> source files referencing it.""" + found: dict[str, set[str]] = {} + for path in _OMNI_ROOT.rglob("*.py"): + try: + tree = ast.parse(path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): # pragma: no cover - defensive + continue + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == _GUARDED: + found.setdefault(node.attr, set()).add(str(path.relative_to(_OMNI_ROOT.parent))) + return found + + +def test_omni_only_uses_prompt_update_helpers_that_exist_upstream() -> None: + referenced = _referenced_attributes() + assert referenced, "expected Omni to reference PromptUpdateDetails somewhere" + + missing = {attr: sorted(files) for attr, files in referenced.items() if not hasattr(PromptUpdateDetails, attr)} + + assert not missing, ( + "Omni calls PromptUpdateDetails helpers that the installed vLLM does not " + f"provide: {missing}. Upstream moved prompt updates to token ids; use " + "select_token_id/select_token_ids with the encoded text." + ) + + +def test_token_oriented_helpers_are_the_ones_upstream_exposes() -> None: + """Pin the replacements, so their removal fails here and not at startup.""" + assert hasattr(PromptUpdateDetails, "select_token_id") + assert hasattr(PromptUpdateDetails, "select_token_ids") + assert not hasattr(PromptUpdateDetails, "select_text"), ( + "select_text is back upstream; re-check the MiniCPM-o replacements." + ) diff --git a/tests/platforms/test_platform_detection.py b/tests/platforms/test_platform_detection.py new file mode 100644 index 000000000..6673a2376 --- /dev/null +++ b/tests/platforms/test_platform_detection.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import logging + +import pytest + +import vllm_omni.platforms as platforms + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_failing_oot_plugin_is_logged_and_falls_back(monkeypatch, caplog): + plugin_name = "raising_oot_plugin" + + def raising_plugin(): + raise RuntimeError("injected platform detection failure") + + monkeypatch.setattr(platforms, "builtin_omni_platform_plugins", {}) + monkeypatch.setattr( + platforms, + "load_omni_plugins_by_group", + lambda _group: {plugin_name: raising_plugin}, + ) + + target_logger = logging.getLogger(platforms.__name__) + target_logger.addHandler(caplog.handler) + previous_level = target_logger.level + target_logger.setLevel(logging.DEBUG) + try: + platform_cls = platforms.resolve_current_omni_platform_cls_qualname() + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(previous_level) + + assert platform_cls == "vllm_omni.platforms.interface.UnspecifiedOmniPlatform" + assert plugin_name in caplog.text + assert "RuntimeError: injected platform detection failure" in caplog.text diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 10f0853e0..3313e691f 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -2975,7 +2975,9 @@ async def benchmark( "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, + "total_input_sequences": metrics.total_input_sequences, "request_throughput": metrics.request_throughput, + "input_sequence_throughput": metrics.input_sequence_throughput, "total_token_throughput": metrics.total_token_throughput, "input_lens": [output.prompt_len for output in outputs], "errors": [output.error for output in outputs], diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index aa71245a0..9123b794d 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -107,6 +107,17 @@ class OmniSchedulerMixin: def _init_omni_io_scheduling_state(self) -> None: """Initialize scheduler state shared by AR and generation stages.""" + # Every omni worker forces the v1 model runner (the v2 runner carries no + # omni hooks), so the scheduler has to agree or the two disagree about + # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where + # 0.28 did not, so the scheduler took its v2 fast path -- the one that + # deliberately omits resumed-request token ids -- while the v1 runner + # still read them, and resuming a request raised ``KeyError: `` + # in ``_update_states``. + if getattr(self, "use_v2_model_runner", False): + logger.warning("OMNI scheduler forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + model_config = self.vllm_config.model_config self.chunk_transfer_adapter = ( OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 5b5b8a328..e903960d6 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -121,6 +121,13 @@ class AttentionBackend(ABC): when it unifies cache layouts across layers. Dense diffusion backends conservatively keep the default ``False``; a paged backend should override this only when its kernel actually follows that layout. + + Since vLLM 0.29 this no longer rides on ``AttentionSpec``. It selects + the block-outermost physical layout pinned on + ``CacheConfig.kv_cache_layout`` -- see + ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose + ``is_block_outermost`` is what upstream reads when it sizes + ``KVCacheTensor`` regions. """ return False diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index ac541808b..4b8ca8f49 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported from vllm_omni.diffusion.distributed.parallel_state import get_sp_group from vllm_omni.diffusion.forward_context import ( get_forward_context, @@ -261,15 +262,19 @@ class Attention(nn.Module): return None dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype # Keep backend layout discovery under the same config context used by - # upstream vLLM's attention-spec collector. + # upstream vLLM's attention-spec collector. vLLM 0.29 moved the + # block-stride decision off the spec and onto the single physical + # ``CacheConfig.kv_cache_layout`` resolved before the cache is built, + # so the backend's preference is enforced there (see + # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than + # carried per layer. with set_current_vllm_config(vllm_config): - indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride() + assert_backend_layout_supported(vllm_config, self.attn_backend) return FullAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=dtype, - indexes_kv_by_block_stride=indexes_kv_by_block_stride, non_causal=not self.causal, ) diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py index 31aaaa46e..d1a79d934 100644 --- a/vllm_omni/diffusion/diffusion_kv/initialization.py +++ b/vllm_omni/diffusion/diffusion_kv/initialization.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -16,6 +16,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config from vllm_omni.platforms import current_omni_platform @@ -28,14 +29,35 @@ def build_native_kv_cache_configs( vllm_config: VllmConfig, worker_specs: list[dict[str, KVCacheSpec]], available_memory: list[int], + *, + indexes_kv_by_block_stride: bool = False, ) -> tuple[list[KVCacheConfig], KVCacheConfig]: - """Build rank-local and Scheduler-native configs using vLLM utilities.""" + """Build rank-local and Scheduler-native configs using vLLM utilities. + + ``indexes_kv_by_block_stride`` carries the native backend requirement that + 0.29 removed from ``AttentionSpec``. Every diffusion backend in tree + declares ``False``, so the default reproduces current behaviour; a caller + that learns otherwise from its workers passes ``True`` here and the + block-outermost layout propagates to every rank-local config. + """ from vllm.v1.core import single_type_kv_cache_manager + # 0.29 sizes every KVCacheTensor region from the resolved physical layout, + # so it must be pinned before get_kv_cache_configs. Idempotent, so callers + # that already resolved it (the control plane) are unaffected. + layout = resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, + ) single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config) worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory) - return worker_configs, generate_scheduler_kv_cache_config(worker_configs) + scheduler_config = generate_scheduler_kv_cache_config(worker_configs) + # Workers only receive a KVCacheConfig, so the resolved layout travels on it + # and is adopted into each worker's own CacheConfig before allocation. + for config in (*worker_configs, scheduler_config): + config.kv_cache_layout = layout.name + return worker_configs, scheduler_config def initialize_diffusion_kv_control_plane( diff --git a/vllm_omni/diffusion/diffusion_kv/layout.py b/vllm_omni/diffusion/diffusion_kv/layout.py new file mode 100644 index 000000000..960df395d --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/layout.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Physical KV cache layout selection for the native Diffusion paged cache. + +vLLM 0.29 removed ``AttentionSpec.indexes_kv_by_block_stride`` and replaced it +with a single physical layout per model, resolved once and recorded on +``CacheConfig.kv_cache_layout``. ``KVCacheLayout.is_block_outermost`` is the +direct successor of the old per-spec flag: upstream derives +``interleaved_block_stride`` from it when it lays out ``KVCacheTensor`` regions +(``vllm/v1/core/kv_cache_utils.py``). + +Diffusion never reaches vLLM's engine-core resolution. Its backends subclass +Omni's own ``AttentionBackend`` ABC, which upstream's +``get_supported_kv_cache_layouts`` never sees, so the layout is pinned here +before ``get_kv_cache_configs`` runs. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.v1.kv_cache_interface import KVCacheSpec, KVCacheTensor, compute_layout_strides +from vllm.v1.kv_cache_layout import KVCacheLayout + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +# Layer-outermost, matching the head of upstream's ``_DEFAULT_LAYOUT_PREFERENCE`` +# and reproducing the ``indexes_kv_by_block_stride=False`` that every dense +# diffusion backend declared before 0.29. +DEFAULT_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.LBNHC +# Block-outermost counterpart, for a paged backend whose kernel really does read +# K/V pages by the runtime block stride. +BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.BLNHC + + +def resolve_diffusion_kv_cache_layout( + vllm_config: VllmConfig, + *, + indexes_kv_by_block_stride: bool = False, +) -> KVCacheLayout: + """Pin the physical KV layout for the Diffusion cache and return it. + + A layout already present on the config wins, mirroring upstream's + ``resolve_kv_cache_layout``; it is validated against the requirement rather + than silently overwritten, because an inconsistent layout reads K/V + from the wrong offsets instead of raising. + """ + + required_block_outermost = bool(indexes_kv_by_block_stride) + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None: + # A config double with nowhere to record the choice; report what the + # requirement implies without pretending it was pinned. + return BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + + current = getattr(cache_config, "kv_cache_layout", None) + if current is not None: + layout = KVCacheLayout[current] + if layout.is_block_outermost != required_block_outermost: + raise ValueError( + "Diffusion KV cache layout " + f"{layout.name} (is_block_outermost=" + f"{layout.is_block_outermost}) contradicts the attention " + "backend, which requires is_block_outermost=" + f"{required_block_outermost}." + ) + return layout + + layout = BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + cache_config.kv_cache_layout = layout.name + return layout + + +def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout: + """Adopt the control plane's resolved layout inside a worker. + + The engine core resolves the layout once and every worker adopts it before + allocating (``CacheConfig.kv_cache_layout``). Diffusion workers receive only + a ``KVCacheConfig``, so the name rides along on ``kv_cache_config`` and is + copied onto the worker's own config here. Without this, anything calling + ``get_resolved_kv_cache_layout`` in a worker -- ``init_kv_cache`` does -- + raises. + """ + + name = getattr(kv_cache_config, "kv_cache_layout", None) + if name is None: + # Older configs, or a standalone harness: fall back to the same default + # the control plane pins so both sides still agree. + return resolve_diffusion_kv_cache_layout(vllm_config) + + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None: + return KVCacheLayout[name] + + current = getattr(cache_config, "kv_cache_layout", None) + if current is None: + cache_config.kv_cache_layout = name + elif current != name: + raise ValueError( + f"Worker KV cache layout {current!r} disagrees with the layout resolved by the control plane ({name!r})." + ) + return KVCacheLayout[name] + + +def build_kv_cache_tensor( + spec: KVCacheSpec, + num_blocks: int, + layers: Sequence[str], + *, + layout: KVCacheLayout | None = None, + offset: int = 0, +) -> KVCacheTensor: + """A ``KVCacheTensor`` covering ``layers`` for ``num_blocks`` blocks. + + Before 0.29 a tensor named the layers sharing it + (``KVCacheTensor(size=..., shared_by=[...])``); 0.29 replaced that with an + explicit layout, so layer ``l``'s block ``b`` starts at + ``offset + l * layer_stride + b * block_stride``. Note this lays the layers + out at distinct offsets rather than aliasing them, which is what upstream + now does for a single cache group. + + The strides come from upstream's own ``compute_layout_strides``, called the + way ``vllm.v1.core.kv_cache_utils`` calls it, because hand-computed strides + read K/V from the wrong offsets silently instead of raising. + """ + + layout = layout or DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + layer_names = list(layers) + # Mirrors kv_cache_utils: the block stride is pinned up front only for + # block-outermost layouts, where every block repeats the same packing. + bytes_per_block = spec.page_size_bytes * len(layer_names) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), + ) + return KVCacheTensor( + size=bytes_per_block * num_blocks, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, + ) + + +def assert_backend_layout_supported(vllm_config: VllmConfig, attn_backend: type | None) -> None: + """Fail loudly when a backend's block-stride need contradicts the layout. + + Skipped while the layout is still unresolved -- worker processes collect + specs before the control plane hands them the resolved name -- so this + enforces the contract wherever it is knowable and never crashes a rank that + has simply not been told yet. + """ + + if attn_backend is None: + return + # Spec discovery runs against lightweight config doubles in tests and against + # worker configs that have not been handed the resolved name yet; in both + # cases there is simply nothing to check. + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None or getattr(cache_config, "kv_cache_layout", None) is None: + return + indexes_by_block_stride = getattr(attn_backend, "indexes_kv_by_block_stride", None) + if indexes_by_block_stride is None: + return + resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=indexes_by_block_stride(), + ) diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py index c5ab3fb16..20f4db165 100644 --- a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py +++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py @@ -18,6 +18,10 @@ from vllm.v1.worker.gpu.block_table import BlockTables from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.layout import ( + adopt_kv_cache_layout, + assert_backend_layout_supported, +) from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, @@ -285,13 +289,16 @@ class DiffusionKVModelRunnerBackend: layers=self._kv_cache_layer_adapters, resolve_row=self._resolve_paged_attention_row, ) + # vLLM 0.29 reads the resolved physical layout when allocating, + # and dropped attn_groups/cache_dtype from the signature. + adopt_kv_cache_layout(self.vllm_config, kv_cache_config) + for layer_adapter in self._kv_cache_layer_adapters.values(): + assert_backend_layout_supported(self.vllm_config, getattr(layer_adapter, "attn_backend", None)) init_kv_cache( kv_caches, self.vllm_config.compilation_config.static_forward_context, kv_cache_config, - attn_groups, self.device, - self.vllm_config.cache_config.cache_dtype, kernel_block_sizes, self.vllm_config, ) diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py index 6a7862669..a76a5434c 100644 --- a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py +++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py @@ -105,10 +105,12 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase): attn_backend, ulysses_degree=ulysses_degree, ) + # vLLM 0.29 carries the block-stride decision on the resolved physical + # ``CacheConfig.kv_cache_layout`` instead of the spec; the backend's + # requirement is enforced against that layout rather than copied here. canonical_spec = replace( spec, num_kv_heads=num_kv_heads, - indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), ) self.layer_name = layer_name self.spec = canonical_spec diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 59d2c9a62..37b2ed835 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -1,3 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from importlib import import_module from importlib.util import find_spec import torch @@ -98,6 +102,7 @@ class RotaryEmbedding(CustomOp): self.interleaved = not is_neox_style self.half_head_dim = half_head_dim self.apply_rotary_emb_flash_attn = None + self.apply_rotary_emb_vllm_flash_attn = None self.has_mindie = False # ``find_spec("flash_attn")`` is True as long as *any* package publishes # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships @@ -111,6 +116,17 @@ class RotaryEmbedding(CustomOp): self.apply_rotary_emb_flash_attn = apply_rotary except ImportError: pass + # CUDA wheels vendor this kernel under vllm_flash_attn, while source + # and CPU-only installations may not contain the generated package. + # Resolve it lazily and retain the native implementation as a portable + # fallback for those environments. + if current_omni_platform.is_cuda(): + try: + self.apply_rotary_emb_vllm_flash_attn = import_module( + "vllm.vllm_flash_attn.layers.rotary" + ).apply_rotary_emb + except ImportError: + pass if find_spec("mindiesd") is not None: self.has_mindie = True @@ -135,12 +151,18 @@ class RotaryEmbedding(CustomOp): cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb - cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin) + if self.apply_rotary_emb_vllm_flash_attn is None: + return apply_rotary_emb_torch( + x, + cos, + sin, + interleaved=self.interleaved, + ) + x, squeezed = _ensure_batch_dim(x) - output = apply_rotary_emb( + output = self.apply_rotary_emb_vllm_flash_attn( x, cos, sin, @@ -231,13 +253,14 @@ class RotaryEmbeddingWan(RotaryEmbedding): cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb + if self.apply_rotary_emb_vllm_flash_attn is None: + return self.forward_native(x, cos, sin) if cos.dim() > 2: cos = cos.reshape(-1, cos.shape[-1]) sin = sin.reshape(-1, sin.shape[-1]) - return apply_rotary_emb( + return self.apply_rotary_emb_vllm_flash_attn( x, cos, sin, diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 6598525da..fcfd48b4f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -454,11 +454,10 @@ class HunyuanImage3Pipeline( # Note: guidance_emb and timestep_r_emb are no longer skipped # to support HunyuanImage-3.0-Distil and MeanFlow distilled models - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) ) - return loader.load_weights(weights) def prepare_seed(self, seed=None, batch_size=1): # random seed diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py index 6b0433bbe..7228fa655 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py @@ -24,7 +24,7 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.batch_invariant import linear_batch_invariant +from vllm.model_executor.determinism.batch_invariant import linear_batch_invariant from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index c82ace9c4..696728b90 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -50,6 +50,17 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) +def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Rotate interleaved pairs before rounding back to the activation dtype. + + Qwen-Image's reference uses complex FP32 multiplication. Rounding the + frequencies to BF16 before rotation loses positional precision; those + errors accumulate across the denoising steps. + """ + paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) + + def _normalize_qwen_image_weight_name(name: str) -> str: name = name.removeprefix("transformer.") if ".to_out.0." in name: @@ -642,15 +653,23 @@ class QwenImageCrossAttention(nn.Module): txt_query = self.norm_added_q(txt_query) txt_key = self.norm_added_k(txt_key) - img_cos = torch.real(vid_freqs).to(img_query.dtype) - img_sin = torch.imag(vid_freqs).to(img_query.dtype) - txt_cos = torch.real(txt_freqs).to(txt_query.dtype) - txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) + if img_query.device.type == "cuda": + img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs) + img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs) + txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs) + txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs) + else: + # Retain the platform-specific kernels on other accelerators, + # which may not support complex tensors. + img_cos = torch.real(vid_freqs).to(img_query.dtype) + img_sin = torch.imag(vid_freqs).to(img_query.dtype) + txt_cos = torch.real(txt_freqs).to(txt_query.dtype) + txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) - img_query = self.rope(img_query, img_cos, img_sin) - img_key = self.rope(img_key, img_cos, img_sin) - txt_query = self.rope(txt_query, txt_cos, txt_sin) - txt_key = self.rope(txt_key, txt_cos, txt_sin) + img_query = self.rope(img_query, img_cos, img_sin) + img_key = self.rope(img_key, img_cos, img_sin) + txt_query = self.rope(txt_query, txt_cos, txt_sin) + txt_key = self.rope(txt_key, txt_cos, txt_sin) seq_len_txt = encoder_hidden_states.shape[1] joint_query = torch.cat([txt_query, img_query], dim=1) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index bd6a71562..44bd9058c 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -443,8 +443,8 @@ class OmniEngineArgs(EngineArgs): @dataclass class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs): @classmethod - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: - parser = AsyncEngineArgs.add_cli_args(parser) + def add_cli_args(cls, parser: argparse.ArgumentParser, async_args_only: bool = False) -> argparse.ArgumentParser: + parser = AsyncEngineArgs.add_cli_args(parser, async_args_only=async_args_only) parser = OmniEngineArgs._add_omni_specific_args(parser) return parser diff --git a/vllm_omni/engine/stage_engine_core_proc_manager.py b/vllm_omni/engine/stage_engine_core_proc_manager.py index b14c62d8f..94a3e58b1 100644 --- a/vllm_omni/engine/stage_engine_core_proc_manager.py +++ b/vllm_omni/engine/stage_engine_core_proc_manager.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Process manager for omni stage engine subprocesses. @@ -83,6 +83,12 @@ class StageEngineCoreProcManager(CoreEngineProcManager): if local_engine_count <= 0: raise ValueError(f"local_engine_count must be > 0, got {local_engine_count}") + # Mirrors the vLLM 0.29 parent __init__: the inherited shutdown() reads + # this to bound how long in-flight requests may drain. Omitting it makes + # shutdown raise AttributeError, which leaves the engine core + # subprocesses alive and hangs interpreter exit until the job timeout. + self._request_shutdown_timeout = vllm_config.shutdown_timeout + context = get_mp_context() common_kwargs: dict[str, object] = { "vllm_config": vllm_config, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 765549bcb..de97f7aa3 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -54,7 +54,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import OmniInputPreprocessor +from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -1549,25 +1549,24 @@ class _TokenOnlyRenderer(BaseRenderer): def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer: - return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None) + return omni_renderer_cls(_TokenOnlyRenderer)(stage_vllm_config, tokenizer=None) def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: - """Build the shared stage-0 input processor.""" + """Build the shared stage-0 input processor. + + The renderer is the Omni subclass of the upstream renderer class that + ``renderer_from_config`` would have picked (or of the token-only renderer + when the stage skips tokenizer initialization), so upstream's + ``InputProcessor`` runs unmodified on top of it. + """ patch_generation_config_if_needed(stage_vllm_config.model_config) if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)): - input_processor = InputProcessor( - vllm_config=stage_vllm_config, - renderer=_build_token_only_renderer(stage_vllm_config), - ) + renderer = _build_token_only_renderer(stage_vllm_config) else: - input_processor = InputProcessor(vllm_config=stage_vllm_config) - input_processor.input_preprocessor = OmniInputPreprocessor( - vllm_config=stage_vllm_config, - renderer=input_processor.renderer, - ) - return input_processor + renderer = build_omni_renderer(stage_vllm_config) + return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer) def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 948f7e53a..0012039e0 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -51,9 +51,9 @@ from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetric from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: - from vllm.inputs.preprocess import InputPreprocessor from vllm.tokenizers import TokenizerLike from vllm.v1.engine import PauseMode + from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType @@ -1714,7 +1714,7 @@ class AsyncOmni(EngineClient, OmniBase): # ==================== EngineClient Interface ==================== - async def get_input_preprocessor(self) -> InputPreprocessor: + async def get_input_preprocessor(self) -> InputProcessor: """Get input preprocessor.""" return self.input_processor @@ -1755,10 +1755,12 @@ class AsyncOmni(EngineClient, OmniBase): """ logger.debug("Weight update start requested (no-op in omni)") - async def finish_weight_update(self) -> None: + async def finish_weight_update(self, weight_version: str | None = None) -> None: """Finish the current weight update. Omni does not currently support weight transfer, so this is a no-op. + ``weight_version`` is accepted for upstream ``EngineClient`` protocol + compatibility (RLHF weight-transfer routers pass it positionally). """ logger.debug("Weight update finish requested (no-op in omni)") diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 78b212f98..5c0204a33 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -18,7 +18,7 @@ from typing import Any, cast import uvloop from vllm.entrypoints.cli.types import CLISubcommand -from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args +from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py index 6e4d0a0e0..f41d7205b 100644 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ b/vllm_omni/entrypoints/duplex/chat_fallback.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator from typing import Any from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm_omni.entrypoints.duplex.protocol import DuplexSession diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d976eb181..e020be7bf 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -34,7 +34,10 @@ from starlette.types import ASGIApp, Receive, Scope, Send from vllm.engine.protocol import EngineClient from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template -from vllm.entrypoints.launcher import serve_http, terminate_if_errored +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.launchers.cli_args import make_arg_parser +from vllm.entrypoints.launchers.launcher import serve_http, terminate_if_errored +from vllm.entrypoints.launchers.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server @@ -43,19 +46,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, ) -from vllm.entrypoints.openai.cli_args import make_arg_parser # yapf conflicts with isort for this block # yapf: disable # yapf: enable from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - ModelCard, - ModelList, - ModelPermission, - RequestResponseMetadata, -) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses @@ -64,14 +59,13 @@ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAISer from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens - -# vLLM < 0.28 keeps create_error_response under serve.utils (it is not -# re-exported from vllm.entrypoints.serve); 0.28+ moved it under -# serve.exception_handling and re-exports it from the package root. -try: - from vllm.entrypoints.serve import create_error_response -except ImportError: - from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response +from vllm.entrypoints.serve.engine.protocol import ( + ErrorResponse, + ModelCard, + ModelList, + ModelPermission, +) # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. @@ -85,7 +79,6 @@ from vllm.entrypoints.serve.utils.api_utils import ( ) from vllm.entrypoints.serve.utils.orca_metrics import metrics_header from vllm.entrypoints.serve.utils.request_logger import RequestLogger -from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime from vllm.entrypoints.speech_to_text.transcription.serving import ( OpenAIServingTranscription, diff --git a/vllm_omni/entrypoints/openai/batch_serving.py b/vllm_omni/entrypoints/openai/batch_serving.py index bc5ef36ca..74d78dff7 100644 --- a/vllm_omni/entrypoints/openai/batch_serving.py +++ b/vllm_omni/entrypoints/openai/batch_serving.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import time @@ -11,7 +11,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionResponse, ChatCompletionResponseChoice, ) -from vllm.entrypoints.openai.engine.protocol import ( +from vllm.entrypoints.serve.engine.protocol import ( ErrorResponse, UsageInfo, ) diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index aad7ea658..86a8ba031 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -8,7 +11,7 @@ from typing import cast from uuid import uuid4 import numpy as np -from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone from vllm.logger import init_logger diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index bfa241a40..f5cb24efd 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -54,8 +54,17 @@ except ImportError: soundfile = None +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionCall, + FunctionDefinition, + RequestResponseMetadata, + ToolCall, +) from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs -from vllm.entrypoints.launcher import terminate_if_errored +from vllm.entrypoints.launchers.launcher import terminate_if_errored from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -69,22 +78,15 @@ from vllm.entrypoints.openai.chat_completion.serving import ( _get_mm_token_counts, _make_prompt_tokens_details, ) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ErrorInfo, - ErrorResponse, - FunctionCall, - FunctionDefinition, - RequestResponseMetadata, - ToolCall, - UsageInfo, -) from vllm.entrypoints.openai.parser.harmony_utils import ( get_streamable_parser_for_assistant, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.serve.engine.protocol import ( + ErrorInfo, + ErrorResponse, + UsageInfo, +) from vllm.entrypoints.serve.engine.typing import ChatLikeRequest from vllm.entrypoints.serve.utils.api_utils import should_include_usage from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls @@ -700,7 +702,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if negative_prompt is not None: tprompt["negative_prompt"] = negative_prompt # Always attach mm_processor_kwargs (possibly empty) so - # OmniInputPreprocessor._process_text routes through the + # OmniRenderer routes through the # multimodal processor path. Without it, the preprocessor # falls back to plain _tokenize_prompt and AR-based image-gen # models like GLM-Image never see their image-generation diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 88d9f090d..a9256d953 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -24,12 +24,10 @@ import soundfile as sf import torch from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing -from vllm.entrypoints.launcher import terminate_if_errored -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - RequestResponseMetadata, -) +from vllm.entrypoints.launchers.launcher import terminate_if_errored +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector from vllm.utils import random_uuid diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 35a39b2ea..0f01fb28f 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -156,7 +156,9 @@ def _decode_video_bytes( frames_array, metadata = loader.load_bytes( video_bytes, num_frames=num_frames, - backend="pyav", + # vLLM 0.29 removed the "pyav" decoder backend; "opencv" is + # upstream's default and the remaining CPU-only option. + backend="opencv", keep=keep, ) except Exception as exc: diff --git a/vllm_omni/entrypoints/openai/video_stream_session.py b/vllm_omni/entrypoints/openai/video_stream_session.py index ead3b1ec5..f77fcf74e 100644 --- a/vllm_omni/entrypoints/openai/video_stream_session.py +++ b/vllm_omni/entrypoints/openai/video_stream_session.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Video streaming session manager and WebSocket handler. Provides ``VideoStreamConfig``, ``VideoStreamSession`` (frame/audio buffer @@ -20,7 +20,7 @@ from fastapi import WebSocket, WebSocketDisconnect from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index f18a00ada..f4cb416f6 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Engine-level KV cache orchestration for AR-Diffusion models. This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a @@ -22,10 +23,10 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus +from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( @@ -134,7 +135,7 @@ def build_kv_manager( """ layer_names = list(layer_names) group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) - tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)] + tensors = [build_kv_cache_tensor(spec, num_blocks, layer_names)] config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=tensors, diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py new file mode 100644 index 000000000..bb02fe961 --- /dev/null +++ b/vllm_omni/inputs/mm_processor.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Multimodal data plumbing for Omni's custom, non-HF processors.""" + +from collections.abc import Mapping +from typing import TypeVar + +from transformers import BatchFeature +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo + +_I = TypeVar("_I", bound=BaseProcessingInfo) + + +class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): + """Share data/passthrough handling, leaving model processing in its owner. + + These processors implement ``_call_hf_processor`` locally, sometimes + without any HF processor object. Prompt tokens are prepared separately + in ``apply``; the upstream base still owns caching and prompt updates. + """ + + _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" + + def _apply_hf_processor_main( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> BatchFeature: + kwargs = dict(hf_processor_mm_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) + 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 3da1ef897..dc4105b4d 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,200 +1,118 @@ -from typing import Any +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni renderer: the upstream renderer class with Omni engine-input routing. -from typing_extensions import assert_never -from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput -from vllm.inputs.preprocess import InputPreprocessor -from vllm.logger import init_logger -from vllm.renderers.inputs import SingletonDictPrompt +``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a +concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream +renderer (tokenizer, executors, multimodal cache and every public entry point +come from upstream unchanged) and only ``_process_singleton`` / +``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves +the concrete class the same way upstream ``renderer_from_config`` does. +""" -from vllm_omni.inputs.data import ( - OmniEmbedsPrompt, - OmniTextPrompt, - OmniTokenInputs, - OmniTokensPrompt, - token_inputs_omni, -) +from __future__ import annotations -logger = init_logger(__name__) +from typing import TYPE_CHECKING, cast + +from vllm.renderers import BaseRenderer +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.tokenizers import TokenizerLike -class OmniInputPreprocessor(InputPreprocessor): - """Input preprocessor for omni models. +class OmniRenderer: + """Mixin that routes Omni prompts through the upstream renderer. - Extends the base InputPreprocessor to handle omni-specific input - types including prompt embeddings and additional information payloads. - Supports processing tokens, embeddings, text, and multimodal inputs. + It has no base class of its own; ``omni_renderer_cls`` prepends it to a + concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that + renderer at runtime (hence the casts below). + + Two extensions over upstream ``_process_singleton``: + + * a token prompt that carries ``mm_processor_kwargs`` but no + ``multi_modal_data`` still goes through the multimodal processor (AR + image-generation models such as GLM-Image read their target size from + those kwargs); + * Omni-only prompt keys (``prompt``, ``cache_salt``, + ``additional_information``, ``model_intermediate_buffer``) are copied + onto the engine input. """ - def _process_text( - self, - parsed_content: OmniTextPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> OmniTokenInputs | MultiModalInput: - """Process text prompts with support for mm_processor_kwargs. + @staticmethod + def _with_omni_extras(inputs, prompt): + for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"): + if key in prompt: + inputs[key] = prompt[key] + return inputs - Extends base class to support mm_processor_kwargs without multi_modal_data. - This is needed for models like GLM-Image where text-to-image generation - requires processor kwargs (target_h, target_w) to format the prompt. - """ - prompt_text = parsed_content["prompt"] - mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {} - # When the deprecated raw-prompt path is used, process_inputs does - # not pass mm_uuids to preprocess(). Fall back to reading it from - # the prompt dict so the Renderer's _validate_mm_uuids can see it. - effective_mm_uuids = mm_uuids or parsed_content.get("multi_modal_uuids") + @staticmethod + 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") - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_text, - multi_modal_data, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - elif "mm_processor_kwargs" in parsed_content: - # Presence — not truthiness. An explicitly-set empty dict still - # signals "route through the multimodal processor" (needed for - # AR-based image-gen where the HF processor supplies its own - # defaults and scaffold). - inputs = self._process_multimodal( - prompt_text, + def _process_singleton(self, prompt, *, skip_mm_cache: bool = False): + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = renderer._process_multimodal( + prompt["prompt_token_ids"], {}, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) else: - prompt_token_ids = self._tokenize_prompt( - prompt_text, - tokenization_kwargs=tokenization_kwargs, + inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) + + async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False): + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = await renderer._process_multimodal_async( + prompt["prompt_token_ids"], + {}, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) - inputs = token_inputs_omni( - prompt_token_ids, - prompt_embeds=parsed_content.get("prompt_embeds"), - additional_information=parsed_content.get("additional_information"), - model_intermediate_buffer=parsed_content.get("model_intermediate_buffer"), - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt - - return inputs - - def _process_tokens( - self, - parsed_content: OmniTokensPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> OmniTokenInputs | MultiModalInput: - prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) - prompt_embeds = parsed_content.get("prompt_embeds") - additional_information = parsed_content.get("additional_information") - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - - multi_modal_data = parsed_content.get("multi_modal_data") - - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data: - inputs = self._process_multimodal( - prompt_token_ids, - multi_modal_data, - parsed_content.get("mm_processor_kwargs"), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=parsed_content.get("multi_modal_uuids"), - ) - else: - inputs = token_inputs_omni( - prompt_token_ids=prompt_token_ids, - prompt_embeds=prompt_embeds, - additional_information=additional_information, - model_intermediate_buffer=model_intermediate_buffer, - ) - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - if additional_information is not None: - inputs["additional_information"] = additional_information - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if prompt_text := parsed_content.get("prompt"): - inputs["prompt"] = prompt_text - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt + inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) - return inputs - def _process_embeds( - self, - parsed_content: OmniEmbedsPrompt, - ) -> EmbedsInput: - """Process embeddings prompt with omni-specific extensions. +_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {} - Extends base _process_embeds to handle additional_information payload - for direct transfer between pipeline stages. - """ - # Call parent implementation for base embeds processing - inputs = super()._process_embeds(parsed_content) - # Add omni-specific additional_information if present - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information # type: ignore[typeddict-unknown-key] - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer # type: ignore[typeddict-unknown-key] +def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]: + """Return the Omni subclass of ``base_cls`` (created once per base class).""" + if issubclass(base_cls, OmniRenderer): + return base_cls + cls = _OMNI_RENDERER_CLASSES.get(base_cls) + if cls is None: + name = f"Omni{base_cls.__name__}" + cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name}) + _OMNI_RENDERER_CLASSES[base_cls] = cls + return cls - return inputs - def _prompt_to_llm_inputs( - self, - prompt: SingletonDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> SingletonInput: - """ - Extract the singleton inputs from a prompt. +def build_omni_renderer( + vllm_config: VllmConfig, + *, + renderer_cls: type[BaseRenderer] | None = None, + tokenizer: TokenizerLike | None = None, +) -> BaseRenderer: + """Build the Omni renderer for ``vllm_config``. - Arguments: - - * prompt: single encoder or decoder input prompt - - Returns: - - * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance - """ - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - - if "prompt_token_ids" in prompt: - return self._process_tokens( - prompt, # type: ignore[arg-type] - ) - - if "prompt" in prompt: - return self._process_text( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) - - assert_never(prompt) # type: ignore[arg-type] + With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``: + the tokenizer and renderer mode come from the model config and the class + from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally + ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages. + """ + model_config = vllm_config.model_config + if renderer_cls is None: + if tokenizer is None: + tokenizer = cached_tokenizer_from_config(model_config) + renderer_mode, *_ = tokenizer_args_from_config(model_config) + renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode) + return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer) diff --git a/vllm_omni/model_executor/models/audex/audex_omni_processing.py b/vllm_omni/model_executor/models/audex/audex_omni_processing.py index 5e4c53908..7b80f53f8 100644 --- a/vllm_omni/model_executor/models/audex/audex_omni_processing.py +++ b/vllm_omni/model_executor/models/audex/audex_omni_processing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): # inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_features.py @@ -33,13 +33,14 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + SOUND_TOKEN = "" SOUND_START_TOKEN = "" SOUND_END_TOKEN = "" @@ -176,7 +177,7 @@ class AudexDummyInputsBuilder(BaseDummyInputsBuilder[AudexProcessingInfo]): return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios, overrides=overrides)} -class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]): +class AudexMultiModalProcessor(OmniMultiModalProcessor[AudexProcessingInfo]): def _call_hf_processor( self, prompt: str, diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 495e0037f..30dced9d0 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from math import isqrt from typing import Any @@ -20,7 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.bagel import BagelForConditionalGeneration from vllm.model_executor.models.interfaces import MultiModalEmbeddings from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -221,6 +222,20 @@ class OmniBagelDataParser(MultiModalDataParser): class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): IMG2IMG_PLACEHOLDER = "<|fim_middle|>" + def apply(self, inputs, timing_ctx): + num_img2img = inputs.mm_data_items.get_all_counts().get("img2img", 0) + if num_img2img: + token_id = self.info.get_tokenizer().get_vocab().get(self.IMG2IMG_PLACEHOLDER) + if token_id is not None: + prompt_ids = list(inputs.prompt) + missing = num_img2img - prompt_ids.count(token_id) + if missing > 0: + inputs = replace( + inputs, + prompt=[token_id] * missing + prompt_ids, + ) + return super().apply(inputs, timing_ctx) + @staticmethod def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]: # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing. @@ -241,69 +256,51 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn "pixel_values_img2img": MultiModalFieldConfig.batched("img2img"), } - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> "BatchFeature": - has_image = "images" in mm_data - has_img2img = "pixel_values_img2img" in mm_data - - if has_img2img and self.IMG2IMG_PLACEHOLDER not in prompt: - prompt = f"{self.IMG2IMG_PLACEHOLDER}{prompt}" - - if has_image and has_img2img: - outputs = BatchFeature() - - img_data = dict(mm_data) - if "pixel_values_img2img" in img_data: - del img_data["pixel_values_img2img"] - kwargs_img = dict(mm_kwargs) - kwargs_img["is_img2img"] = False - out_img = super()._call_hf_processor(prompt, img_data, kwargs_img, tok_kwargs) - if "pixel_values" in out_img: - outputs["pixel_values"] = out_img["pixel_values"] - for k, v in out_img.items(): - if k != "pixel_values": - outputs[k] = v - - img2img_data = dict(mm_data) - if "images" in img2img_data: - del img2img_data["images"] - img2img_data["images"] = img2img_data.pop("pixel_values_img2img") - kwargs_img2img = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) - kwargs_img2img["is_img2img"] = True - out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) - if "pixel_values" in out_img2img: - outputs["pixel_values_img2img"] = out_img2img["pixel_values"] - for k, v in out_img2img.items(): - if k not in outputs: - outputs[k] = v - - return outputs - - elif has_img2img: - mm_data = dict(mm_data) - mm_data["images"] = mm_data.pop("pixel_values_img2img") - mm_kwargs = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) - mm_kwargs["is_img2img"] = True - outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) - if "pixel_values" in outputs: - outputs["pixel_values_img2img"] = outputs.pop("pixel_values") - return outputs - - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) - - def _hf_processor_applies_updates( - self, - prompt_text: str, mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + ) -> 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) + prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + processor = self.info.get_hf_processor(**hf_processor_mm_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_outputs = self.info.ctx.call_hf_processor( + processor, + {"text": prompt_text, "images": mm_data["images"]}, + image_kwargs, + ) + processed_data.update(image_outputs) + + if has_img2img: + img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs) + img2img_kwargs["is_img2img"] = True + img2img_outputs = self.info.ctx.call_hf_processor( + processor, + { + "text": prompt_text, + "images": mm_data["pixel_values_img2img"], + }, + img2img_kwargs, + ) + pixel_values = img2img_outputs.pop("pixel_values", None) + if pixel_values is not None: + processed_data["pixel_values_img2img"] = pixel_values + for key, value in img2img_outputs.items(): + processed_data.setdefault(key, value) + + if not has_image and not has_img2img: + processed_data = BatchFeature(dict(passthrough_data)) + else: + processed_data.update(passthrough_data) + return processed_data def _get_prompt_updates( self, @@ -372,7 +369,8 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn embed_mask = [True] * num_vae_total + [False] + [True] * num_vit_total return PromptUpdateDetails( full=tokens, - is_embed=lambda _tok, _seq, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), + # vLLM 0.29 calls is_embed with just the full token list. + is_embed=lambda _full, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), ) replacements.append( @@ -1102,11 +1100,12 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): filtered_weights.append((mapped_name, tensor)) loader = AutoWeightsLoader( - self, - skip_prefixes=["vit_pos_embed.pos_embed"], - ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."], + self, ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."] + ) + loaded = loader.load_weights( + filtered_weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"vit_pos_embed.pos_embed": None}), ) - loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) loaded |= self._load_moe_gen_weights(moe_gen_weights) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 84652cada..558e96b65 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from functools import partial from math import gcd from threading import Lock @@ -24,7 +25,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -38,6 +38,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.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -108,7 +109,45 @@ class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): ) -class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): +class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): + def apply(self, inputs: ProcessorInputs, timing_ctx): + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) + + text_token, text_token_len = extract_text_token( + prompt_text, + self.tokenizer, + config.allowed_special, + ) + if inputs.mm_data_items.get_all_counts().get("audio", 0): + reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text") + if not isinstance(reference_text, str): + raise ValueError(f"prompt text is None : {reference_text}") + prompt_text_token, prompt_text_token_len = extract_text_token( + reference_text, + self.tokenizer, + config.allowed_special, + ) + text_token, _ = concat_text_with_prompt_ids( + text_token, + text_token_len, + prompt_text_token, + prompt_text_token_len, + ) + + inputs = replace( + inputs, + prompt=text_token.reshape(-1).tolist(), + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None: cached_model_dir = getattr(self, "_cached_model_dir", None) if cached_model_dir == model_dir: @@ -376,15 +415,6 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal "embedding": MultiModalFieldConfig.batched("audio"), } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py index efbe7b0d6..108bd36c5 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 Tencent. from collections.abc import Iterable, Mapping, Sequence from typing import Any @@ -17,7 +20,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, @@ -26,6 +28,7 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights @@ -106,10 +109,7 @@ class CovoAudioDummyInputsBuilder(BaseDummyInputsBuilder[CovoAudioProcessingInfo ) -class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingInfo]): - def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool: - return False - +class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]): def _call_hf_processor( self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object] ) -> BatchFeature: diff --git a/vllm_omni/model_executor/models/gepard/gepard_talker.py b/vllm_omni/model_executor/models/gepard/gepard_talker.py index 02f5da533..37c1e66a9 100644 --- a/vllm_omni/model_executor/models/gepard/gepard_talker.py +++ b/vllm_omni/model_executor/models/gepard/gepard_talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Gepard-1.0 native-AR talker. Single-stage AR TTS on a vLLM-native Qwen3.5 backbone. Each step samples one @@ -22,7 +22,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM -from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig @@ -199,8 +199,10 @@ class GepardTalkerForConditionalGeneration(nn.Module): # NanoCodec loads lazily on first decode, so load_format=dummy and # NeMo-less environments never pay for it. - loader = AutoWeightsLoader(self, skip_prefixes=["mtp.", "ref_compressor."]) - return loader.load_weights(iter(rest)) + loader = AutoWeightsLoader(self) + return loader.load_weights( + iter(rest), mapper=WeightsMapper(orig_to_new_prefix={"mtp.": None, "ref_compressor.": None}) + ) def _get_or_create_state(self, request_id: str) -> _GepardState: st = self._active_states.get(request_id) 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 77aef398b..e97a78c25 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 @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm_image/modeling_glm_image.py @@ -24,6 +24,7 @@ import math import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from functools import lru_cache from typing import Annotated, Literal @@ -81,7 +82,6 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, ) from vllm.multimodal.parse import ( - ImageProcessorItems, MultiModalDataItems, MultiModalDataParser, ) @@ -341,6 +341,30 @@ class GlmImageDummyInputsBuilder(BaseDummyInputsBuilder[GlmImageProcessingInfo]) } +_GLM_IMAGE_GRID_FACTOR = 32 + + +def _build_target_shape_scaffold(processor, *, height: int, width: int, is_text_to_image: bool) -> str: + """Return the grid scaffold HF appends to a GLM-Image generation prompt. + + Mirrors ``GlmImageProcessor._build_prompt_with_target_shape`` (identical in + transformers 5.13 and 5.14) using only the processor's public token + attributes, so this does not depend on a private HF method that a stub or + a future release may not provide. + """ + factor = _GLM_IMAGE_GRID_FACTOR + token_h = (height // factor * factor) // factor + token_w = (width // factor * factor) // factor + grid_bos, grid_eos, bos = processor.grid_bos_token, processor.grid_eos_token, processor.bos_token + scaffold = f"{grid_bos}{token_h} {token_w}{grid_eos}" + if is_text_to_image: + ratio = token_h / token_w + prev_token_h = int(math.sqrt(ratio) * (factor // 2)) + prev_token_w = int(math.sqrt(1 / ratio) * (factor // 2)) + scaffold += f"{grid_bos}{prev_token_h} {prev_token_w}{grid_eos}" + return scaffold + bos + + class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo]): """ Multimodal processor for GLM-Image. @@ -358,265 +382,38 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo return self._apply_hf_processor(inputs, timing_ctx) return super()._cached_apply_hf_processor(inputs, timing_ctx) - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Call the HuggingFace processor. - - For text-to-image mode (no images), we need to: - 1. Build the prompt with target grid dimensions - 2. Build the image_grid_thw tensor for M-RoPE position encoding - - For image-to-image mode: - 1. Process source images through the image processor - 2. Build prompt with image placeholders expanded - 3. Build image_grid_thw including source and target grids - """ - processor = self.info.get_hf_processor() - - # Get target dimensions from mm_kwargs or use defaults - target_h = mm_kwargs.get("target_h", 1024) if mm_kwargs else 1024 - target_w = mm_kwargs.get("target_w", 1024) if mm_kwargs else 1024 - - logger.debug( - f"_call_hf_processor: target dimensions for generation: {target_h}x{target_w}, mm_kwargs={mm_kwargs}" - ) - - if not mm_data or not mm_data.get("images"): - # Text-to-image mode - if processor is not None: - # Build messages format expected by processor - messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] - - # Use apply_chat_template which handles target dimensions - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", - ) - - return hf_inputs - else: - # Fallback: just tokenize (this won't work properly for generation) - tokenizer = self.info.get_tokenizer() - prompt_ids = tokenizer.encode(prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - # Image-to-image mode - # NOTE: Use "images" (plural) - this is what vLLM's ImageProcessorItems.get_processor_data() returns - images = mm_data.get("images") - if not isinstance(images, list): - images = [images] - - logger.debug( - f"_call_hf_processor i2i: num_images={len(images)}, image_types={[type(img).__name__ for img in images]}" - ) - - if processor is not None: - # Build messages with image objects directly in content - # This is how GlmImageProcessor expects images - embedded in the content dict - # NOT as a separate images= parameter - # - # IMPORTANT: Remove <|image|> placeholders from prompt since apply_chat_template - # will automatically insert them for each image in content. Having both leads to - # index out of bounds when processing image_grid_thw. - clean_prompt = prompt.replace("<|image|>", "") - content = [] - for img in images: - content.append({"type": "image", "image": img}) - content.append({"type": "text", "text": clean_prompt}) - messages = [{"role": "user", "content": content}] - - logger.debug(f"_call_hf_processor: calling apply_chat_template with {len(images)} images in content") - - # Use apply_chat_template - processor will process images when they're in content - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", + def apply(self, inputs, timing_ctx): + """Add the generation scaffold before upstream expands source images.""" + num_images = inputs.mm_data_items.get_all_counts().get("image", 0) + prompt_ids = list(inputs.prompt) + if num_images: + image_token_id = getattr( + self.info.get_hf_config(), + "image_token_id", + 167855, ) + missing = num_images - prompt_ids.count(image_token_id) + if missing > 0: + prompt_ids = [image_token_id] * missing + prompt_ids - logger.debug(f"_call_hf_processor: apply_chat_template returned keys: {list(hf_inputs.keys())}") - - # IMPORTANT (i2i): vLLM multimodal encoder must see source-only grids - # (matching pixel_values and number of images), but M-RoPE needs full - # grids (source + target) to compute correct decode positions. - image_grid_thw = hf_inputs.get("image_grid_thw") - if image_grid_thw is not None: - # Preserve full grids for M-RoPE. - hf_inputs["mrope_image_grid_thw"] = image_grid_thw - - # Expose source-only grids for MM. - # In most i2i requests, we process one prompt at a time here, - # so `len(images)` is the number of source images. - num_source_images = len(images) - if image_grid_thw.shape[0] != num_source_images: - source_grids = image_grid_thw[:num_source_images] - hf_inputs["image_grid_thw"] = source_grids - logger.debug( - "_call_hf_processor: adjusted image_grid_thw for MM from %s to %s (num_source_images=%d)", - tuple(image_grid_thw.shape), - tuple(source_grids.shape), - num_source_images, - ) - - # Debug: Analyze input_ids for image tokens - input_ids = hf_inputs.get("input_ids") - if input_ids is not None: - if hasattr(input_ids, "tolist"): - ids_list = input_ids.tolist() - if isinstance(ids_list[0], list): - ids_list = ids_list[0] # Unbatch - else: - ids_list = list(input_ids) - - # Get image token ID from config - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - - # Count image tokens - image_token_count = ids_list.count(image_token_id) - logger.debug( - f"_call_hf_processor: input_ids length={len(ids_list)}, " - f"image_token_id={image_token_id}, " - f"image_token_count={image_token_count}" - ) - - # Log first/last few tokens to understand structure - logger.debug(f"_call_hf_processor: first 20 tokens: {ids_list[:20]}") - logger.debug(f"_call_hf_processor: last 20 tokens: {ids_list[-20:]}") - - # Find positions of image tokens - image_positions = [i for i, t in enumerate(ids_list) if t == image_token_id] - if image_positions: - logger.debug(f"_call_hf_processor: image token positions (first 10): {image_positions[:10]}") - - return hf_inputs - else: - # Fallback without processor - this is not ideal but prevents crashes - logger.warning("GlmImageProcessor not available, using fallback for i2i") - tokenizer = self.info.get_tokenizer() - hf_config = self.info.get_hf_config() - - # Get image token - image_token_id = getattr(hf_config, "image_token_id", 167855) - try: - image_token = tokenizer.convert_ids_to_tokens(image_token_id) - except Exception: - image_token = "<|image|>" - - # Build prompt with image placeholders - image_placeholders = image_token * len(images) - full_prompt = f"{image_placeholders}{prompt}" - prompt_ids = tokenizer.encode(full_prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - 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. - - GLM-Image requires special handling because apply_chat_template always - adds a target <|image|> placeholder in addition to source image placeholders. - This causes an IndexError when the HF processor tries to find grid info - for the target placeholder (which doesn't exist for source-only processing). - - Solution: Call the image processor directly to get pixel_values and - image_grid_thw, bypassing apply_chat_template's target handling. - """ - mm_counts = mm_items.get_all_counts() - num_images = mm_counts.get("image", 0) - - if num_images == 0: - # No images - call parent implementation - return super()._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # For i2i mode, we need to process images directly with the image processor - # to avoid the apply_chat_template target placeholder issue processor = self.info.get_hf_processor() - image_processor = processor.image_processor - - # Get images from mm_items - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - - logger.debug(f"_apply_hf_processor_mm_only: processing {len(image_list)} images directly") - - # Process images directly with image processor - image_inputs = image_processor( - images=image_list, - return_tensors="pt", - ) - - # Get grid info for source images only (no target) - pixel_values = image_inputs.get("pixel_values") - image_grid_thw = image_inputs.get("image_grid_thw") - if image_grid_thw is not None and image_grid_thw.shape[0] != num_images: - # Be defensive: some processors may include extra target grids. - image_grid_thw = image_grid_thw[:num_images] - image_inputs["image_grid_thw"] = image_grid_thw - - logger.debug( - f"_apply_hf_processor_mm_only: pixel_values shape=\ - {pixel_values.shape if pixel_values is not None else None}, " - f"image_grid_thw shape={image_grid_thw.shape if image_grid_thw is not None else None}" - ) - - # Build input_ids with image token placeholders - # The _get_prompt_updates returns PromptReplacement(target=[image_token_id], ...) - # which needs to find image tokens in input_ids to replace them. - # We need to include one image_token_id per image so the replacement can work. tokenizer = self.info.get_tokenizer() - image_token_id = tokenizer.convert_tokens_to_ids("<|image|>") - - # Build input_ids: [image_token] * num_images + tokenized text - # This way _apply_prompt_updates can find the image tokens and replace them - dummy_text = self.dummy_inputs.get_dummy_text(mm_counts) - text_ids = tokenizer.encode(dummy_text, add_special_tokens=False) - input_ids = [image_token_id] * num_images + text_ids - - logger.debug( - f"_apply_hf_processor_mm_only: built input_ids with {num_images} image tokens + {len(text_ids)} text tokens" + # Append HF's target-shape scaffold without decoding/re-tokenizing the + # user's tokens. Text-to-image needs both target and preview grids; i2i + # needs only the target. Grid metadata alone does not supply these AR + # tokens. + target_grid = self._build_generation_grids(inputs.hf_processor_mm_kwargs)[0] + suffix = _build_target_shape_scaffold( + processor, + height=int(target_grid[1]) * 32, + width=int(target_grid[2]) * 32, + is_text_to_image=num_images == 0, ) - - return BatchFeature( - dict( - input_ids=[input_ids], - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - ), - tensor_type="pt", - ) - - def _apply_hf_processor_text_only( - self, prompt_text: str, hf_processor_mm_kwargs: Mapping[str, object], tokenization_kwargs: Mapping[str, object] - ) -> list[int]: - prompt_ids, _, _ = super()._apply_hf_processor_text_mm( - prompt_text=prompt_text, - mm_items=MultiModalDataItems({}), - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - return prompt_ids + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + if prompt_ids[-len(suffix_ids) :] != suffix_ids: + prompt_ids.extend(suffix_ids) + inputs = replace(inputs, prompt=prompt_ids) + return super().apply(inputs, timing_ctx) def _build_generation_grids(self, hf_processor_mm_kwargs: Mapping[str, object]) -> torch.Tensor: """Build generation grids for M-RoPE decode positions. @@ -661,171 +458,39 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: - """ - Override to handle GLM-Image i2i mode correctly. - - Problem: When vLLM processes cached mm items (enable_hf_prompt_update=False), - the base implementation: - 1. Gets prompt_ids from _apply_hf_processor_text_only (no image tokens) - 2. Gets mm_data from _apply_hf_processor_mm_only - 3. Returns is_update_applied=False - - This causes _apply_prompt_updates to fail because prompt_ids has no image tokens. - - Solution: For i2i mode, we build prompt_ids that include image placeholders, - and return is_update_applied=False so _apply_prompt_updates can expand them. - """ + ) -> 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)}) - logger.debug(f"_apply_hf_processor_main: mm_counts={mm_counts}, num_images={num_images}") + 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) + images = processor_data.get("images") + if not images: + return BatchFeature(dict(passthrough_data)) - if num_images == 0 and isinstance(prompt, str): - # t2i mode or normal flow - use parent implementation - prompt_ids = self._apply_hf_processor_text_only( - prompt_text=prompt, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # t2i has no source images, so mm features cannot provide image_grid_thw. - # Provide explicit generation grids for M-RoPE to avoid fallback token parsing - # (which can degrade high-resolution spatial positions, e.g. 1920x1920). - try: - mrope_grid_thw = self._build_generation_grids(hf_processor_mm_kwargs) - mm_processed_data["mrope_image_grid_thw"] = mrope_grid_thw - logger.debug( - "_apply_hf_processor_main t2i: mrope_image_grid_thw=%s", - mrope_grid_thw.tolist(), - ) - except Exception as e: - logger.warning("_apply_hf_processor_main t2i: failed to set mrope_image_grid_thw: %s", e) - - return prompt_ids, mm_processed_data, False - - # i2i mode: use unified HF processor path only. - # This avoids drift between duplicated manual/HF i2i implementations. - logger.debug( - "_apply_hf_processor_main: i2i mode (enable_hf_prompt_update=%s), num_images=%s", - enable_hf_prompt_update, - num_images, + processor = self.info.get_hf_processor() + image_inputs = processor.image_processor( + images=list(images), + return_tensors="pt", ) + image_grid_thw = image_inputs.get("image_grid_thw") + if image_grid_thw is not None: + source_grid_thw = image_grid_thw[:num_images] + image_inputs["image_grid_thw"] = source_grid_thw - if not isinstance(prompt, str): - # Online OpenAI chat preprocessing can arrive here with tokenized - # prompts (list[int]) before serving_chat replaces engine prompt - # with the clean text prompt. Do not fail the whole request. - logger.warning( - "_apply_hf_processor_main i2i: got tokenized prompt type=%s; " - "using compatibility path for preprocessing", - type(prompt).__name__, + target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype) + image_inputs["mrope_image_grid_thw"] = torch.cat( + [source_grid_thw, target_grid], + dim=0, ) - prompt_ids = list(prompt) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # Preserve full grids for M-RoPE decode (source + target), while - # keeping image_grid_thw source-only for MM batching. - try: - source_grid_thw = mm_processed_data.get("image_grid_thw") - if source_grid_thw is not None and isinstance(source_grid_thw, torch.Tensor): - target_h = ( - hf_processor_mm_kwargs.get("target_h") - if isinstance(hf_processor_mm_kwargs.get("target_h"), int) - else None - ) - target_w = ( - hf_processor_mm_kwargs.get("target_w") - if isinstance(hf_processor_mm_kwargs.get("target_w"), int) - else None - ) - if target_h is None or target_w is None: - target_h = ( - hf_processor_mm_kwargs.get("height") - if isinstance(hf_processor_mm_kwargs.get("height"), int) - else 1024 - ) - target_w = ( - hf_processor_mm_kwargs.get("width") - if isinstance(hf_processor_mm_kwargs.get("width"), int) - else 1024 - ) - - factor = 32 - token_h = max(1, target_h // factor) - token_w = max(1, target_w // factor) - target_grid = torch.tensor([[1, token_h, token_w]], dtype=source_grid_thw.dtype) - mm_processed_data["mrope_image_grid_thw"] = torch.cat([source_grid_thw, target_grid], dim=0) - except Exception: - pass - - # Prompt updates will expand image placeholders in this compatibility path. - return prompt_ids, mm_processed_data, False - - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - if not image_list: - raise ValueError("GLM-Image i2i requires at least one source image in mm_items") - - hf_inputs = self._call_hf_processor( - prompt=prompt, - mm_data={"images": image_list}, - mm_kwargs=hf_processor_mm_kwargs, - tok_kwargs=tokenization_kwargs, - ) - - input_ids = hf_inputs.get("input_ids") - if input_ids is None: - raise ValueError("HF i2i processor returned no input_ids") - - if isinstance(input_ids, torch.Tensor): - prompt_ids = input_ids[0].tolist() if input_ids.dim() > 1 else input_ids.tolist() - else: - prompt_ids = ( - input_ids[0] - if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list) - else list(input_ids) - ) - - mm_processed_data = BatchFeature(dict(), tensor_type="pt") - for key in ("pixel_values", "image_grid_thw", "mrope_image_grid_thw"): - value = hf_inputs.get(key) - if value is not None: - mm_processed_data[key] = value - - image_grid_thw = mm_processed_data.get("image_grid_thw") - mrope_grid_thw = mm_processed_data.get("mrope_image_grid_thw") - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - image_token_count = prompt_ids.count(image_token_id) - logger.debug( - "_apply_hf_processor_main i2i(HF): num_images=%s, prompt_len=%s, image_token_count=%s, " - "source_grid_shape=%s, mrope_grid_shape=%s", - num_images, - len(prompt_ids), - image_token_count, - tuple(image_grid_thw.shape) if image_grid_thw is not None else None, - tuple(mrope_grid_thw.shape) if mrope_grid_thw is not None else None, - ) - - # HF processor already expanded image placeholders in input_ids. - return prompt_ids, mm_processed_data, True + image_inputs.update(passthrough_data) + return image_inputs def _get_mm_fields_config( self, @@ -836,7 +501,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo Get the multimodal field configuration. For GLM-Image i2i mode: - - image_grid_thw has been sliced in _call_hf_processor to only include source images + - image_grid_thw contains only source grids from _apply_hf_processor_main - pixel_values has shape [total_patches, C, H, W] - only for source images For t2i mode: @@ -844,7 +509,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ result = {} - # Get image_grid_thw if present (already sliced in _call_hf_processor) + # Get the source-only image_grid_thw if present. image_grid_thw = hf_inputs.get("image_grid_thw") if "pixel_values" in hf_inputs and image_grid_thw is not None: @@ -861,45 +526,21 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo result["pixel_values"] = MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes) - # Register image_grid_thw - it's been sliced in _call_hf_processor + # Register image_grid_thw - it's been sliced in _apply_hf_processor_main # to only include source image grids, so batching will work correctly result["image_grid_thw"] = MultiModalFieldConfig.batched("image") + if "mrope_image_grid_thw" in hf_inputs: + result["mrope_image_grid_thw"] = MultiModalFieldConfig.shared( + "image", + num_source_images, + keep_on_cpu=True, + ) + logger.debug(f"_get_mm_fields_config: result keys: {list(result.keys())}") return result - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - """ - Return whether the HF processor applies prompt updates. - - For GLM-Image i2i mode, the HF processor's apply_chat_template already - expands <|image|> to N tokens (e.g., 4096 for 64x64 grid). - - By returning True, we tell vLLM that HF processor DID apply prompt updates, - so vLLM will use _find_mm_placeholders to locate the expanded tokens - instead of trying to apply replacements. - - For t2i mode (no images), there are no image placeholders to expand. - """ - # Check if we have images (i2i mode) - num_images = mm_items.get_all_counts().get("image", 0) - if num_images > 0: - logger.debug( - f"_hf_processor_applies_updates: returning True for i2i mode " - f"(num_images={num_images}) - HF processor already expanded tokens" - ) - return True - - # For t2i mode (no images), use default behavior - return True - def _get_prompt_updates( self, mm_items: MultiModalDataItems, @@ -909,17 +550,8 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ Get prompt updates for image tokens. - For GLM-Image image-to-image mode, the HF processor's apply_chat_template - already expands each <|image|> placeholder to the correct number of - image tokens (grid_h * grid_w tokens per source image). - - The HF processor does: - 1. Replace each <|image|> with num_image_tokens copies of <|placeholder|> - 2. Replace all <|placeholder|> back to <|image|> - - So the tokenized input already has the expanded tokens. We use - target=[image_token_id] to match each occurrence of the image token, - similar to how Qwen2VL handles this pattern. + Expand each source-image placeholder into grid_h * grid_w image tokens + through vLLM's token-based prompt updates. We use image_grid_thw from out_mm_kwargs to get the actual processed grid size, following the Qwen2VL pattern. This is critical because the HF processor @@ -930,8 +562,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo """ hf_config = self.info.get_hf_config() - # Get image token ID - this is the token that appears multiple times - # in the tokenized input after HF processor expansion + # Match the source-image placeholder in the tokenized input. image_token_id = getattr(hf_config, "image_token_id", 167855) # Debug: log mm_items info 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 fb3517036..e98c227a4 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -39,7 +39,6 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -51,6 +50,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.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 @@ -364,7 +364,7 @@ class GLMTTSMultiModalProcessingInfo(BaseProcessingInfo): ) -class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): +class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): """GLM-TTS voice-clone processor. Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs @@ -373,6 +373,53 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi handoff. """ + def apply(self, inputs: ProcessorInputs, timing_ctx): + source_tokenizer = self.info.get_tokenizer() + prompt_text = source_tokenizer.decode( + inputs.prompt, + skip_special_tokens=False, + ) + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components( + model_dir, + config, + load_voice_clone=False, + ) + + normalized_text = _normalize_glm_tts_processor_text( + self.text_frontend, + prompt_text, + ) + text_ids = self._encode_text(normalized_text) + prompt_parts = [] + if inputs.mm_data_items.get_all_counts().get("audio", 0): + reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text") + if not isinstance(reference_text, str) or not reference_text.strip(): + raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].") + normalized_reference = _normalize_glm_tts_processor_text( + self.text_frontend, + reference_text, + add_trailing_space=True, + ) + prompt_parts.append(self._encode_text(normalized_reference)) + prompt_parts.extend( + [ + text_ids, + torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long), + ] + ) + prompt_ids = torch.cat(prompt_parts, dim=1).reshape(-1).tolist() + inputs = replace( + inputs, + prompt=prompt_ids, + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components( self, model_dir: str, @@ -599,15 +646,6 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi if key in hf_inputs } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _cached_apply_hf_processor(self, inputs: ProcessorInputs, timing_ctx: Any): # GLM-TTS builds the actual AR prompt from both the request text and # mm_processor_kwargs["prompt_text"]. The base cache path separates text diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 8fbac95f7..738e469f1 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -1669,8 +1669,10 @@ class _BackboneWrapper(nn.Module): self.config = config def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - from vllm.model_executor.models.utils import AutoWeightsLoader + from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None - loader = AutoWeightsLoader(self, skip_prefixes=skip) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip or ())}) + ) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py new file mode 100644 index 000000000..cd1546505 --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Vendored from vllm/model_executor/models/hunyuan_v1.py at vllm@2cf0a6915c +# (the upstream base vllm-omni origin/main was rebased against). Upstream +# removed the native HunYuan V1 implementation in #53272 / #53615 (migrated +# to the Transformers modeling backend). vllm-omni's HunyuanImage3 model +# subclasses the removed HunYuanModel / HunYuanSparseMoeBlock / HunYuanMLP, +# so those internals are preserved here to keep the omni diffusion model +# working against current upstream. +# Weight loading, expert mapping, and QKV splitting are owned by Omni's +# HunyuanModel subclass; their unused upstream implementations are omitted. +"""Inference-only HunYuan V1 internals vendored for vllm-omni HunyuanImage3.""" + +from itertools import islice + +import torch +from torch import nn +from transformers import PretrainedConfig +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import EagleModelMixin +from vllm.model_executor.models.utils import ( + PPMissingLayer, + make_layers, +) +from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backend import AttentionType + + +def _is_moe(config: PretrainedConfig) -> bool: + num_experts = getattr(config, "num_experts", None) + if isinstance(num_experts, int): + return num_experts > 1 + if isinstance(num_experts, list) and num_experts: + # Ensure all elements are integers before calling max. + if all(isinstance(e, int) for e in num_experts): + return max(num_experts) > 1 + else: + return False + return False + + +def _get_cla_factor(config: PretrainedConfig) -> int: + if not getattr(config, "use_cla", False): + return 1 + return getattr(config, "cla_share_factor", 1) + + +class HunYuanMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + reduce_results=reduce_results, + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class HunYuanAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + if hasattr(config, "head_dim") and config.head_dim: + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=True, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + ori_k = k + if self.use_qk_norm: + q = self.query_layernorm( + q.view(-1, self.num_heads, self.head_dim), + ) + k = self.key_layernorm( + k.view(-1, self.num_kv_heads, self.head_dim), + ) + + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output, (ori_k, v) + + +class HunYuanCrossAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + if hasattr(config, "head_dim"): + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.q_proj = ColumnParallelLinear( + hidden_size, + hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=True, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + attn_type=AttentionType.ENCODER_DECODER, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + ) -> torch.Tensor: + assert kv_states is not None + ori_k, v = kv_states # use last layer kv, + k = ori_k + q, _ = self.q_proj(hidden_states) + k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding + q, _ = self.rotary_emb(positions, q, k_tmp) + if self.use_qk_norm: + q = self.query_layernorm( + q.view(-1, self.num_heads, self.head_dim), + ) + k = self.key_layernorm( + k.view(-1, self.num_kv_heads, self.head_dim), + ) + + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output, (ori_k, v) + + +class HunYuanSparseMoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + layer_id: int = -1, + prefix: str = "", + enable_eplb: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + + self.ep_group = get_ep_group().device_group + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.num_experts + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}." + ) + + # Get layer_id topk if config.moe_topk is a list + if isinstance(config.moe_topk, list): + assert layer_id >= 0 + assert len(config.moe_topk) > layer_id + top_k = config.moe_topk[layer_id] + else: + top_k = config.moe_topk + + # If it is moe, moe_intermediate_size is preferred + intermediate_size = config.intermediate_size + if config.moe_intermediate_size is not None: + intermediate_size = ( + config.moe_intermediate_size + if isinstance(config.moe_intermediate_size, int) + else config.moe_intermediate_size[layer_id] + ) + + # Load balancing settings. + vllm_config = get_current_vllm_config() + eplb_config = vllm_config.parallel_config.eplb_config + self.enable_eplb = enable_eplb + + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + if config.use_mixed_mlp_moe > 0: + # Get layer_id num_shared_expert if config.num_shared_expert is + # a list. + if isinstance(config.num_shared_expert, list): + assert layer_id >= 0 + assert len(config.num_shared_expert) > layer_id + num_shared_expert = config.num_shared_expert[layer_id] + else: + num_shared_expert = config.num_shared_expert + + self.shared_mlp = HunYuanMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size * num_shared_expert, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_mlp", + ) + else: + self.shared_mlp = None + + self.experts = FusedMoEFactory( + shared_experts=self.shared_mlp, + num_experts=self.n_routed_experts, + top_k=top_k, + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + renormalize=top_k > 1, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + + return final_hidden_states.view(orig_shape) + + +class HunYuanDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_id: int = -1, + enable_eplb: bool = False, + ) -> None: + super().__init__() + assert layer_id >= 0 + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.intermediate_size = ( + config.intermediate_size + if isinstance(config.intermediate_size, int) + else config.intermediate_size[layer_id] + ) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False) + cla_factor = _get_cla_factor(config) + attention_type = ( + AttentionType.ENCODER_DECODER if layer_id >= 0 and layer_id % cla_factor != 0 else AttentionType.DECODER + ) + if attention_type == AttentionType.DECODER: + self.self_attn = HunYuanAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_id=layer_id, + ) + elif attention_type == AttentionType.ENCODER_DECODER: + self.self_attn = HunYuanCrossAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_id=layer_id, + ) + else: + raise RuntimeError(f"Unsupported attention type: {attention_type}") + + if _is_moe(config): + self.mlp = HunYuanSparseMoeBlock( + config=config, + quant_config=quant_config, + layer_id=layer_id, + prefix=f"{prefix}.mlp", + enable_eplb=enable_eplb, + ) + else: + self.mlp = HunYuanMLP( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + kv_states: tuple[torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states, ori_kv_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + kv_states=kv_states, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual, ori_kv_states + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + # positions is of shape (xd, seq_len) if xdrope is enabled for hunyuan-vl, + # otherwise (seq_len, ). + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class HunYuanModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + eplb_config = vllm_config.parallel_config.eplb_config + enable_eplb = vllm_config.parallel_config.enable_eplb + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.config = config + self.quant_config = quant_config + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + quant_config=quant_config, + ) + else: + self.embed_tokens = PPMissingLayer() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: HunYuanDecoderLayer( + config=config, + layer_id=int(prefix.split(".")[-1]), + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + enable_eplb=enable_eplb, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + cla_factor = _get_cla_factor(self.config) + prev_kv_states = None + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for i, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)): + hidden_states, residual, kv_states = layer( + positions, + hidden_states, + residual, + prev_kv_states, + ) + + if getattr(self.config, "use_cla", False) and i % cla_factor == 0: + prev_kv_states = kv_states + else: + prev_kv_states = None + + self._maybe_add_hidden_state(aux_hidden_states, i + 1, hidden_states, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 58b67d8aa..b15a19982 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import gc import math import typing @@ -45,13 +45,6 @@ from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) -from vllm.model_executor.models.hunyuan_v1 import ( - HunYuanMLP, - HunYuanModel, - HunYuanSparseMoeBlock, - _get_cla_factor, - _is_moe, -) from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsLoRA, @@ -63,6 +56,7 @@ from vllm.model_executor.models.interfaces import ( from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, _merge_multimodal_embeddings, is_pp_missing_parameter, maybe_prefix, @@ -78,7 +72,6 @@ from vllm.multimodal.parse import ( ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, @@ -91,6 +84,14 @@ 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.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import ( + HunYuanMLP, + HunYuanModel, + HunYuanSparseMoeBlock, + _get_cla_factor, + _is_moe, +) from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer @@ -1027,7 +1028,7 @@ class HunyuanImage3DummyInputsBuilder(BaseDummyInputsBuilder[HunyuanImage3Proces } -class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3ProcessingInfo]): +class HunyuanImage3MultiModalProcessor(OmniMultiModalProcessor[HunyuanImage3ProcessingInfo]): """Multimodal processor for HunyuanImage3 model.""" def _call_hf_processor( @@ -1046,15 +1047,6 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long) return batch_feature - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, @@ -2180,12 +2172,11 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo "timestep_r_emb", ] skip_prefixes.extend(unexpected_keywords) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) + loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(weights) + loaded_params = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) + ) return loaded_params def get_language_model(self) -> torch.nn.Module: 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 4309d6177..4dbc9a455 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2025 Xiaomi Corporation. import os from collections.abc import Iterable, Mapping, Sequence @@ -27,14 +30,13 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) -from vllm.multimodal.utils import group_mm_kwargs_by_modality +from vllm.multimodal.utils import group_and_batch_mm_kwargs from vllm.sequence import IntermediateTensors from vllm.utils.cache import LRUCache from vllm.utils.collection_utils import is_list_of @@ -42,6 +44,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.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import ( NO_INTERLEAVE_NEXT_TOKEN_ID, @@ -348,7 +351,7 @@ class MiMoAudioDataParser(MultiModalDataParser): return AudioProcessorItems(new_audios) -class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProcessingInfo]): +class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProcessingInfo]): def _call_hf_processor( self, prompt: str, @@ -484,7 +487,7 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_mimo_audio, ) ] @@ -626,7 +629,7 @@ class MiMoAudioForConditionalGeneration( ) mm_kwargs.append((mm_feature.modality, mm_item)) - for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality( + for modality, num_items, mm_kwargs_group in group_and_batch_mm_kwargs( mm_kwargs, device=self.device, pin_memory=self.pin_memory, 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 c90adb467..bb01e522c 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 @@ -390,35 +390,25 @@ class MimoAudioDummyInputsBuilder(BaseDummyInputsBuilder[MimoAudioProcessingInfo class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingInfo]): - def _call_hf_processor( + def _preprocess_hf_mm_data( self, - prompt: str, mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, Any], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: + 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 - - # Text-only input not supported in composite processor - if not mm_data.get("audio", []): - prompt_ids = self.info.get_tokenizer().encode(prompt) - prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - feature_extractor = self.info.get_feature_extractor(**mm_kwargs) - mm_kwargs = dict( - **mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, - ) - - return super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) + 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 def _get_mm_fields_config( self, @@ -480,7 +470,7 @@ class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingIn return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_mimo_audio, ) ] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index d7e31beb3..d34cce465 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright (c) Ant Group. All rights reserved. # Adapted from: @@ -21,7 +22,7 @@ from transformers import AutoTokenizer, Qwen2Config, Qwen2Model from transformers.utils.hub import cached_file from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin @@ -509,12 +510,13 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) if self._standalone: weights = self._iter_talker_safetensors() - loader = AutoWeightsLoader( - self, - skip_prefixes=["audio_vae."], # loaded separately - skip_substrs=["rotary_embed.inv_freq"], # non-persistent buffer + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=WeightsMapper( + orig_to_new_prefix={"audio_vae.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None} + ), ) - loaded = loader.load_weights(weights) logger.info("Loaded %d talker weights from checkpoint", len(loaded)) if self.audio_vae is not None and self._vae_weight_source is not None: 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 03cf87abb..1338f2606 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 @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright 2024 ANT Group and the HuggingFace Inc. team. # Adapted from Ming repository modeling_bailingmm2.py and processing_bailingmm2.py @@ -8,6 +9,7 @@ import os from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import replace from typing import Annotated, Any import numpy as np @@ -388,7 +390,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO updates.append( PromptReplacement( modality="image", - target=PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_image, ) ) @@ -396,7 +401,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO updates.append( PromptReplacement( modality="video", - target=PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_video, ) ) @@ -404,7 +412,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO updates.append( PromptReplacement( modality="audio", - target=PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_audio, ) ) @@ -451,90 +462,71 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO return config - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + def apply(self, inputs, timing_ctx): + """Add Ming image-generation query tokens before prompt updates.""" + modalities = inputs.hf_processor_mm_kwargs.get("modalities") or [] + if "image" in modalities or "img2img" in modalities: + from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + DEFAULT_NUM_QUERY_TOKENS, + IMAGE_PATCH_TOKEN, + ) + + tokenizer = self.info.get_tokenizer() + prompt_ids = list(inputs.prompt) + if "img2img" in modalities: + prompt_ids = ( + tokenizer.encode( + PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + add_special_tokens=False, + ) + + prompt_ids + ) + + image_patch_ids = tokenizer.encode( + IMAGE_PATCH_TOKEN, + add_special_tokens=False, + ) + if not image_patch_ids or not any( + prompt_ids[index : index + len(image_patch_ids)] == image_patch_ids + for index in range(len(prompt_ids) - len(image_patch_ids) + 1) + ): + image_gen_config = getattr( + self.info.ctx.model_config.hf_config, + "image_gen_config", + None, + ) + num_query_tokens = int( + getattr( + image_gen_config, + "num_query_tokens", + DEFAULT_NUM_QUERY_TOKENS, + ) + ) + suffix = ( + tokenizer.encode("", add_special_tokens=False) + + image_patch_ids * num_query_tokens + + tokenizer.encode("", add_special_tokens=False) + ) + prompt_ids.extend(suffix) + + inputs = replace(inputs, prompt=prompt_ids) + + return super().apply(inputs, timing_ctx) def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ): - """Apply Ming-specific prompt rewrites before delegating to vllm's - default text / multimodal processing pipeline. - - This is the single point in the base MM-processor call graph where we - still have BOTH the full prompt and the request-level - ``hf_processor_mm_kwargs`` together. Downstream, vllm's MM caching - design splits the work into - :meth:`_apply_hf_processor_text_only` (prompt, empty kwargs) and - :meth:`_apply_hf_processor_mm_only` (dummy text, real kwargs), so a - per-request flag like ``modalities`` cannot reach the real - tokenization from inside :meth:`_call_hf_processor`. - - For Ming-flash-omni-2.0, image-generation requests carry - ``mm_processor_kwargs["modalities"]`` set to ``["image"]`` (t2i) or - ``["img2img"]`` (image edit). We expand the prompt here with the - ``*N`` query-token block so the expanded - form flows through both the tokenization and (eventual) MM paths. - For img2img, we also prepend ```` so the thinker-side prompt - replacement can locate the reference-image placeholder. - """ - modalities = hf_processor_mm_kwargs.get("modalities") or [] - is_image_gen = "image" in modalities or "img2img" in modalities - if isinstance(prompt, str) and is_image_gen: - from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( - DEFAULT_NUM_QUERY_TOKENS, - maybe_expand_image_gen_prompt, - ) - - # img2img: prepend ```` so the ref-image placeholder survives - # MM caching (when the ref image is cache-warm, the thinker sees 0 - # missing items and cannot inject the placeholder itself). - if "img2img" in modalities: - prompt = "" + prompt - - ig = getattr(self.info.ctx.model_config.hf_config, "image_gen_config", None) - num_query_tokens = getattr(ig, "num_query_tokens", DEFAULT_NUM_QUERY_TOKENS) - prompt = maybe_expand_image_gen_prompt(prompt, num_query_tokens=int(num_query_tokens)) - - return super()._apply_hf_processor_main( - prompt=prompt, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, - ) - - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> BatchFeature: - """Call sub-processors for multimodal inputs and tokenize. + """Call Ming's modality processors without expanding prompt text. We call the image/audio sub-processors directly (instead of going through `MingFlashOmniProcessor.__call__`) so that the high-level - placeholder tokens remain **unexpanded** in the tokenized output. - - NOTE: Image-gen prompt expansion is handled in - :meth:`_apply_hf_processor_main` — by the time we get here, vllm's - MM caching has already split text/mm processing and ``mm_kwargs`` - is empty on the text-only branch. + 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) hf_processor = self.info.get_hf_processor() - tokenizer = self.info.get_tokenizer() data: dict[str, object] = {} @@ -577,11 +569,9 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO ) data.update(audio_outputs) - # Tokenize text with placeholders still intact - text_outputs = tokenizer(prompt, return_tensors="pt", **tok_kwargs) - data.update(text_outputs) - - return BatchFeature(data=data) + processed_data = BatchFeature(data=data) + processed_data.update(passthrough_data) + return processed_data @MULTIMODAL_REGISTRY.register_processor( 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 9dfff5589..315ae159d 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 @@ -21,9 +21,8 @@ import inspect import math import os import warnings -from collections import defaultdict from collections.abc import Callable, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -3403,84 +3402,46 @@ class MiniCPMOMultiModalDataParser(MultiModalDataParser): class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45OmniLLMProcessingInfo]): """Multimodal processor for MiniCPM-o thinker stage.""" - def _apply_hf_processor_main( - self, - prompt: str | list[int], - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: - """ - MiniCPM-O reimplements this to avoid calling HF processor with text-only - when enable_hf_prompt_update=False. The MiniCPM processor asserts - len(image_tags) == len(image_sizes) and fails if given placeholder text - without corresponding image data (e.g. during profiling/cache-miss path). - """ - use_tts = hf_processor_mm_kwargs.get("use_tts", False) - hf_processor_mm_kwargs = {k: v for k, v in hf_processor_mm_kwargs.items() if k != "use_tts"} - if isinstance(prompt, str): - if use_tts: - prompt = prompt + self._TTS_SUFFIX - if enable_hf_prompt_update: - return self._apply_hf_processor_text_mm( - prompt_text=prompt, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - tokenizer = self.info.get_tokenizer() - prompt_ids = _encode_tokens(tokenizer, prompt) - else: - prompt_ids = self._apply_hf_processor_tokens_only(prompt) - if use_tts: - tokenizer = self.info.get_tokenizer() - tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"]) - prompt_ids = list(prompt_ids) + tts_ids - - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return prompt_ids, mm_processed_data, False - _TTS_SUFFIX = "<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>" - def _call_hf_processor( + def apply(self, inputs, timing_ctx): + hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs + if hf_processor_mm_kwargs.get("use_tts", False): + tokenizer = self.info.get_tokenizer() + tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"]) + inputs = replace( + inputs, + prompt=[*inputs.prompt, *tts_ids], + hf_processor_mm_kwargs={ + key: value for key, value in hf_processor_mm_kwargs.items() if key != "use_tts" + }, + ) + return super().apply(inputs, timing_ctx) + + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], + mm_items: MultiModalDataItems, + hf_processor_mm_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) + tokenizer = self.info.get_tokenizer() - - use_tts = mm_kwargs.get("use_tts", False) - mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"} - if use_tts: - prompt = prompt + self._TTS_SUFFIX - - input_ids = torch.tensor([tokenizer.encode(prompt, **tok_kwargs)]) - mm_inputs = self.process_mm_inputs(mm_data, mm_kwargs, tok_kwargs) - - return BatchFeature( + 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) + processed_data = BatchFeature( { "input_ids": input_ids, **mm_inputs, } ) - - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + processed_data.update(passthrough_data) + return processed_data def get_image_prompt_texts( self, @@ -3528,7 +3489,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (audios := mm_data.get("audios")) is None: return {} @@ -3540,11 +3500,10 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn if isinstance(parsed_audios, MiniCPMOAudioEmbeddingItems): audio_inputs = {} else: - audio_inputs = self._base_call_hf_processor( + audio_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.audio_pattern] * len(parsed_audios), mm_data={"audios": [[audio] for audio in parsed_audios]}, mm_kwargs={**mm_kwargs, "chunk_input": True}, - tok_kwargs=tok_kwargs, out_keys={"audio_features", "audio_feature_lens"}, ) @@ -3569,7 +3528,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (images := mm_data.get("images")) is None: return {} @@ -3581,11 +3539,10 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn if isinstance(parsed_images, MiniCPMVImageEmbeddingItems): image_inputs = {} else: - image_inputs = self._base_call_hf_processor( + image_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.image_pattern] * len(parsed_images), mm_data={"images": [[image] for image in parsed_images]}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) return image_inputs @@ -3594,7 +3551,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (videos := mm_data.get("videos")) is None: return {} @@ -3606,14 +3562,13 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn if isinstance(parsed_videos, MiniCPMVVideoEmbeddingItems): video_inputs = {} else: - video_inputs = self._base_call_hf_processor( + video_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.image_pattern * len(video) for video in parsed_videos], mm_data={"images": list(parsed_videos)}, mm_kwargs={ **mm_kwargs, "max_slice_nums": self.info.get_video_max_slice_num(), }, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) @@ -3625,48 +3580,30 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: return { - **self.process_images(mm_data, mm_kwargs, tok_kwargs), - **self.process_videos(mm_data, mm_kwargs, tok_kwargs), - **self.process_audios(mm_data, mm_kwargs, tok_kwargs), + **self.process_images(mm_data, mm_kwargs), + **self.process_videos(mm_data, mm_kwargs), + **self.process_audios(mm_data, mm_kwargs), } - def _base_call_hf_processor( + def _call_hf_processor_on_prompts( self, prompts: list[str], mm_data: Mapping[str, Sequence[object]], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], *, out_keys: set[str], ) -> dict[str, NestedTensors]: - mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"} - # This processor supports zipping prompt and mm_data together - if self.info.get_model_version() in {(2, 6), (4, 0), (4, 5)}: - inputs = super()._call_hf_processor( - prompt=prompts, # type: ignore - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - else: - inputs = defaultdict[str, list[torch.Tensor]](list) + from vllm.model_executor.models.minicpmv import MiniCPMVMultiModalProcessor - for i, prompt in enumerate(prompts): - inputs_one = super()._call_hf_processor( - prompt=prompt, - mm_data={k: v[i] for k, v in mm_data.items()}, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - for k, v in inputs_one.items(): - assert len(v) == 1, (k, len(v)) - inputs[k].append(v[0]) - - return {k: inputs[k] for k in out_keys} + return MiniCPMVMultiModalProcessor._call_hf_processor_on_prompts( + self, + prompts, + mm_data, + {key: value for key, value in mm_kwargs.items() if key != "use_tts"}, + out_keys=out_keys, + ) def _get_prompt_updates( self, @@ -3688,6 +3625,17 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn additional_placeholders.append((modality, sub_pattern)) placeholders += additional_placeholders + # vLLM 0.29 removed PromptUpdateDetails.select_text: prompt updates are + # token oriented now, so encode the replacement text and select the unk + # placeholder positions by token id instead of by text. + unk_token_id = tokenizer.convert_tokens_to_ids("") + + def _select_unk_positions(text: str) -> PromptUpdateDetails: + return PromptUpdateDetails.select_token_id( + tokenizer.encode(text, add_special_tokens=False), + unk_token_id, + ) + image_max_slice_nums = hf_processor_mm_kwargs.get("max_slice_nums") image_use_image_id = hf_processor_mm_kwargs.get("use_image_id") @@ -3696,14 +3644,13 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn image_size = images.get_image_size(item_idx) - return PromptUpdateDetails.select_text( + return _select_unk_positions( self.get_image_prompt_texts( image_size, item_idx, max_slice_nums=None if image_max_slice_nums is None else int(image_max_slice_nums), # type: ignore[arg-type] use_image_id=None if image_use_image_id is None else bool(image_use_image_id), - ), - "", + ) ) def get_video_replacement(item_idx: int): @@ -3712,17 +3659,18 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn frame_size = videos.get_frame_size(item_idx) num_frames = videos.get_num_frames(item_idx) - return PromptUpdateDetails.select_text( - self.get_video_prompt_texts(frame_size, num_frames), - "", - ) + return _select_unk_positions(self.get_video_prompt_texts(frame_size, num_frames)) get_replacement = { "image": get_image_replacement, "video": get_video_replacement, } base_updates = [ - PromptReplacement(modality=modality, target=pattern, replacement=get_replacement[modality]) + PromptReplacement( + modality=modality, + target=tokenizer.encode(pattern, add_special_tokens=False), + replacement=get_replacement[modality], + ) for modality, pattern in placeholders ] @@ -3738,16 +3686,16 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn else: audio_len = audios.get_audio_length(item_idx) - return PromptUpdateDetails.select_text( - self.get_audio_prompt_texts(audio_len), - "", - ) + return _select_unk_positions(self.get_audio_prompt_texts(audio_len)) return [ *base_updates, PromptReplacement( modality="audio", - target=audio_placeholder, + target=tokenizer.encode( + audio_placeholder, + add_special_tokens=False, + ), replacement=get_audio_replacement, ), ] 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 85583f243..b1933af95 100644 --- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -12,6 +12,7 @@ import regex as re import torch import torch.nn as nn from vllm.config import VllmConfig +from vllm.inputs import mm_input from vllm.model_executor.models.qwen3_vl import ( Qwen3VLDummyInputsBuilder, Qwen3VLForConditionalGeneration, @@ -26,10 +27,6 @@ from vllm_omni.model_executor.models.minimax_h3.conditioning import ( MINIMAX_H3_PRESENTATION_TASK_KEY, ) from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - IMAGE_PAD, - VIDEO_PAD, - VISION_END, - VISION_START, build_minimax_h3_presentation, ) from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -96,58 +93,57 @@ class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor): def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: Any, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, ): - task_value = hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY) - if task_value is None: - return super()._apply_hf_processor_main( - prompt, - mm_items, - hf_processor_mm_kwargs, - tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, - ) - if not isinstance(prompt, str): - raise ValueError("MiniMax H3 presentation requires the original prompt text") - - task = str(task_value) - condition_labels = self._condition_labels(hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY)) base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs) - image_count = sum(kind == "image" for kind, _ in condition_labels) - video_count = sum(kind == "video" for kind, _ in condition_labels) - image_slot = f"{VISION_START}{IMAGE_PAD}{VISION_END}" - hf_processor = self.info.get_hf_processor(**base_kwargs) - video_slot = ( - VIDEO_PAD if self._expands_only_video_token(hf_processor) else f"{VISION_START}{VIDEO_PAD}{VISION_END}" - ) - processing_prompt = image_slot * image_count + video_slot * video_count - if not processing_prompt: - processing_prompt = prompt - - _, processed_data, _ = super()._apply_hf_processor_main( - processing_prompt, + return super()._apply_hf_processor_main( mm_items, base_kwargs, - tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, ) + + def apply(self, inputs: Any, timing_ctx: Any): + task_value = inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY) + if task_value is None: + return super().apply(inputs, timing_ctx) + + prompt_ids = self._postprocess_prompt(inputs.prompt) + mm_info = self._cached_apply_hf_processor(inputs, timing_ctx) + base_kwargs = self._base_processor_kwargs(inputs.hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False) + condition_labels = self._condition_labels(inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY)) + out_mm_data = mm_info.kwargs.get_data() image_processor = self.info.get_image_processor(**base_kwargs) ids, _ = _build_minimax_h3_presentation( - self.info.get_tokenizer(), - prompt=prompt, - task=task, + tokenizer, + prompt=prompt_text, + task=str(task_value), condition_labels=condition_labels, - image_grid_thw=processed_data.get("image_grid_thw"), - video_grid_thw=processed_data.get("video_grid_thw"), - video_timestamps=processed_data.get("timestamps"), + image_grid_thw=out_mm_data.get("image_grid_thw"), + video_grid_thw=out_mm_data.get("video_grid_thw"), + video_timestamps=out_mm_data.get("timestamps"), merge_size=int(image_processor.merge_size), ) - return ids.tolist(), processed_data, True + prompt_ids = ids.tolist() + + mm_item_counts = inputs.mm_data_items.get_all_counts() + self._validate_mm_kwargs(mm_info.kwargs, mm_item_counts) + self._validate_mm_updates(mm_info.prompt_updates, mm_item_counts) + mm_placeholders = self._find_mm_placeholders( + prompt_ids, + mm_info.prompt_updates, + ) + self._validate_mm_placeholders(mm_placeholders, mm_item_counts) + mm_placeholder_ranges = { + modality: [item.to_range() for item in placeholders] for modality, placeholders in mm_placeholders.items() + } + return mm_input( + prompt_token_ids=prompt_ids, + mm_kwargs=mm_info.kwargs, + mm_hashes=mm_info.hashes, + mm_placeholders=mm_placeholder_ranges, + ) def _get_prompt_updates( self, diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py index c4fb876fa..9baa7315c 100644 --- a/vllm_omni/model_executor/models/minimax_music3/talker.py +++ b/vllm_omni/model_executor/models/minimax_music3/talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-0 talker for MiniMax Music 3. A Qwen3 backbone predicts one audio frame per decode step. Each frame is @@ -37,7 +37,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.qwen3 import Qwen3Model -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.v1.outputs import SamplerOutput from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -980,8 +980,10 @@ class MiniMaxMusic3TalkerForConditionalGeneration(nn.Module): through vLLM. The audio embedding table and the depth decoder live in a separate component of the repo and are loaded explicitly. """ - loader = AutoWeightsLoader(self, skip_prefixes=["audio_embeddings.", "rvq_decoder."]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={"audio_embeddings.": None, "rvq_decoder.": None}) + ) loaded |= load_depth_decoder_weights( self.vllm_config, audio_embeddings=self.audio_embeddings, diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index 26a6f6b59..a0b0d725e 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -11,6 +11,7 @@ from __future__ import annotations import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace import numpy as np import torch @@ -28,7 +29,6 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -37,6 +37,7 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors +from vllm_omni.inputs.mm_processor import 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 @@ -63,13 +64,29 @@ class OmniVoiceMultiModalProcessingInfo(BaseProcessingInfo): ) -class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): +class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): """Processes text + optional reference audio for OmniVoice. For voice cloning: text + reference audio → tokenized reference For auto voice: text only """ + def apply(self, inputs: ProcessorInputs, timing_ctx): + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) + # Tokenize once here; _apply_hf_processor_main below only encodes the + # reference audio, so the media path never re-tokenizes the text. + prompt_ids = self._encode_prompt_text(prompt_text, inputs.hf_processor_mm_kwargs).reshape(-1).tolist() + inputs = replace( + inputs, + prompt=prompt_ids, + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None: cached_model_dir = getattr(self, "_cached_model_dir", None) if cached_model_dir == model_dir: @@ -97,23 +114,39 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr self._cached_model_dir = model_dir - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: + """Encode the reference audio, if any; the prompt was tokenized in apply(). + + Upstream calls this for every request, including text-only ones and + ones whose audio is already cached (empty ``mm_items``), so it must not + 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) + audio = self._get_reference_audio(mm_data) + if audio is None: + return BatchFeature(dict(passthrough)) + processed = self._encode_reference_audio(*audio) + processed.update(passthrough) + return processed + + def _get_reference_audio(self, mm_data: Mapping[str, object]) -> tuple[object, int] | None: + audio = mm_data.get("audio", None) + if audio is None: + audios = mm_data.get("audios") + if audios is not None: + audio = audios[0], self.info.ctx.get_hf_config().sample_rate + return audio + + def _encode_prompt_text(self, prompt: str, mm_kwargs: Mapping[str, object]) -> torch.Tensor: config = self.info.ctx.get_hf_config() model_dir = self.info.ctx.model_config.model self._ensure_cached_runtime_components(model_dir, config) - audio = mm_data.get("audio", None) - if audio is None: - audio = mm_data.get("audios") - if audio is not None: - audio = audio[0], config.sample_rate - # Build text prompt with control tokens lang = mm_kwargs.get("lang", None) instruct = mm_kwargs.get("instruct", None) @@ -136,19 +169,14 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr full_text = prompt text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>" - text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] + return self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] - if audio is None: - # Text-only path (auto voice mode) - return BatchFeature( - { - "input_ids": text_tokens, - "input_len": [len(text_tokens)], - } - ) + def _encode_reference_audio(self, audio_signal: object, sr: int) -> BatchFeature: + """Voice cloning: encode the reference audio to 8-codebook tokens.""" + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) - # Voice cloning: encode reference audio to tokens - audio_signal, sr = audio if isinstance(audio_signal, np.ndarray): audio_signal = torch.from_numpy(audio_signal).float() if audio_signal.dim() == 1: @@ -169,14 +197,12 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr if ref_audio_tokens.dim() == 3: ref_audio_tokens = ref_audio_tokens.squeeze(0) # [8, T_ref] - ft = BatchFeature( + return BatchFeature( { - "input_ids": text_tokens, "ref_audio_tokens": ref_audio_tokens, # [8, T_ref] "ref_audio_len": [ref_audio_tokens.shape[1]], } ) - return ft def _get_mm_fields_config( self, @@ -188,15 +214,6 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr "ref_audio_len": MultiModalFieldConfig.batched("audio"), } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 35d988895..b3e916f25 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable from functools import cached_property @@ -139,11 +142,11 @@ class Qwen2_5OmniTalkerForConditionalGeneration( return self.language_model.sample(logits, sampling_metadata) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "token2wav."], + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "token2wav.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: total_bytes = 0 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 ba073aad2..6bbefceab 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 @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides.""" import hashlib @@ -9,6 +12,7 @@ import numpy as np import PIL.Image as PILImage import torch from torch import nn +from transformers.feature_extraction_utils import BatchFeature from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniThinkerConfig, ) @@ -360,25 +364,40 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ): """Override to fix use_audio_in_video detection when mm cache returns None.""" - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ): - mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs) - mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(mm_data, mm_kwargs) - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> BatchFeature: + """Apply omni-specific HF kwargs before the upstream Qwen2.5-Omni + processor runs. + + Upstream's ``_apply_hf_processor_main`` is fully custom and no longer + routes through the generic ``_preprocess_hf_mm_data`` / + ``_postprocess_hf_mm_data`` hooks, so the omni video pre-sampling and + 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) + + hf_processor_mm_kwargs = _presampled_videos_hf_kwargs( + processor_data, + hf_processor_mm_kwargs, ) - per_video_mask = mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) + hf_processor_mm_kwargs = _coerce_use_audio_in_video_for_hf_processor( + processor_data, + hf_processor_mm_kwargs, + ) + + processed_data = super()._apply_hf_processor_main( + mm_items, + hf_processor_mm_kwargs, + ) + + per_video_mask = hf_processor_mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) if per_video_mask is not None: - hf_inputs["use_audio_in_video"] = torch.tensor(per_video_mask) - return hf_inputs + processed_data["use_audio_in_video"] = torch.tensor(per_video_mask) + return processed_data def _get_mm_fields_config( self, @@ -530,17 +549,17 @@ class Qwen2_5OmniThinkerMultiModalProcessor( return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_qwen2_audio, ), PromptReplacement( modality="image", - target=image_token, + target=[image_token_id], replacement=partial(get_replacement_qwen2_vision, modality="image"), ), PromptReplacement( modality="video", - target=video_token, + target=[video_token_id], replacement=get_replacement_qwen2_video, ), ] @@ -620,7 +639,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + ) -> MultiModalProcessingInfo: """Cache-aware: processes video/audio pairs as units, using pair-specific cache keys. Falls back to standard cache if no pairs.""" @@ -673,16 +692,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ) with timing_ctx.record("apply_hf_processor"): - ( - prompt_ids, - mm_missing_processed_data, - is_update_applied, - ) = self._apply_hf_processor_main( - prompt=inputs.prompt, + mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=inputs.tokenization_kwargs, - enable_hf_prompt_update=False, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( @@ -714,7 +726,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor( prompt_updates=mm_prompt_updates, ) - return prompt_ids, mm_info, is_update_applied + return mm_info def _derive_audio_from_video_placeholders( self, @@ -785,7 +797,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor( prompt_ids: list[int], mm_kwargs: MultiModalKwargsItems, mm_prompt_updates: MultiModalPromptUpdates, - is_update_applied: bool, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) @@ -794,37 +805,27 @@ class Qwen2_5OmniThinkerMultiModalProcessor( video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) use_audio_in_video = any(video_use_audio_in_video) - if is_update_applied: - mm_placeholders = self._find_mm_placeholders( + if use_audio_in_video and "audio" in mm_prompt_updates: + filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} + prompt_ids, mm_placeholders = self._apply_prompt_updates( + prompt_ids, + filtered_updates, + ) + mm_placeholders = self._derive_audio_from_video_placeholders( + mm_placeholders, + mm_prompt_updates, + video_use_audio_in_video, + ) + else: + prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, mm_prompt_updates, ) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - else: - if use_audio_in_video and "audio" in mm_prompt_updates: - filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - filtered_updates, - ) - mm_placeholders = self._derive_audio_from_video_placeholders( - mm_placeholders, - mm_prompt_updates, - video_use_audio_in_video, - ) - else: - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, - ) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + ) return prompt_ids, mm_placeholders @@ -1397,11 +1398,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration( if self.visual is None: skip_prefixes.extend(["visual."]) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) + | WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index 91eafd2a8..e9e92c904 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable import torch @@ -20,6 +23,7 @@ from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, @@ -441,8 +445,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): return next_tokens def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, + mapper=WeightsMapper(orig_to_new_prefix={"lm_head.": None} if self.config.tie_word_embeddings else {}), ) - return loader.load_weights(weights) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 5c3129be9..67352eecb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The Qwen team. """Inference-only Qwen3-Omni-Moe Code2Wav model.""" @@ -336,11 +336,11 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from HuggingFace checkpoint.""" - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "talker."], # Already loaded above + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "talker.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 8337f8dad..ea3e2fcfe 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable import torch @@ -299,13 +302,12 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( to vLLM's internal structure. Code predictor weights are routed to its custom loader for vocab extension support. """ - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "code2wav."], - # "code_predictor."], - ) + loader = AutoWeightsLoader(self) # Don't apply mapper again since we already did it - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "code2wav.": None}), + ) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index bbfea498b..b2a7d7abe 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The Qwen team. # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. @@ -31,10 +31,7 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F -from packaging.version import Version from transformers import PretrainedConfig -from transformers import __version__ as TRANSFORMERS_VERSION -from transformers.feature_extraction_utils import BatchFeature from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeConfig, Qwen3OmniMoeThinkerConfig, @@ -95,6 +92,9 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import ( from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoderLayer as _Qwen3OmniMoeAudioEncoderLayer, ) +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as _UpstreamQwen3Processor, +) from vllm.model_executor.models.qwen3_omni_moe_thinker import ( _get_feat_extract_output_lengths, ) @@ -124,7 +124,6 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerMultiModalProcessor, _get_request_video_use_audio_in_video, _get_video_second_per_grid_t, - _presampled_videos_hf_kwargs, ) from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, @@ -281,7 +280,12 @@ class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention): self.scaling = self.head_dim**-0.5 - self.qkv = QKVParallelLinear( + # vLLM 0.29 renamed the audio attention projection to ``qkv_proj`` and + # its stacked mapping rewrites checkpoint ``q_proj``/``k_proj``/ + # ``v_proj`` onto that name, so the module must match or weight loading + # fails with "no module or parameter named ...self_attn.qkv_proj". + # The inherited forward also calls ``self.qkv_proj``. + self.qkv_proj = QKVParallelLinear( hidden_size=self.embed_dim, head_size=self.head_dim, total_num_heads=self.num_heads, @@ -289,7 +293,7 @@ class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention): bias=True, disable_tp=disable_tp, quant_config=quant_config, - prefix=f"{prefix}.qkv", + prefix=f"{prefix}.qkv_proj", ) self.out_proj = RowParallelLinear( @@ -696,151 +700,10 @@ Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder class Qwen3OmniMoeThinkerMultiModalProcessor( Qwen2_5OmniThinkerMultiModalProcessor, + _UpstreamQwen3Processor, ): - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_data = dict(mm_data) - mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs)) - audios = mm_data.pop("audios", []) - - def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: - length = x.shape[-1] - if length % hop_length != 0: - pad_length = hop_length - (length % hop_length) - x = np.pad(x, (0, pad_length), mode="constant", constant_values=0) - return x - - # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - feature_extractor = self.info.get_feature_extractor(**mm_kwargs) - hop_length = feature_extractor.hop_length - if audios: - # NOTE: Qwen3-Omni processor accept "audio" - # To make sure the cache works with padding=True, we pre-padded - # the audio to multiple of hop_length. - mm_data["audio"] = [ - pad_to_hop_length(audio, hop_length) - if isinstance(audio, np.ndarray) - else (pad_to_hop_length(audio[0], hop_length), audio[1]) - for audio in audios - ] - - # TODO(Isotr0py): Remove this patch after upstream fix PR - # released and Transformers version update: - # https://github.com/huggingface/transformers/pull/41473 - mm_kwargs = dict(mm_kwargs) - tok_kwargs = dict(tok_kwargs) - mm_kwargs["audio_kwargs"] = dict(mm_kwargs.get("audio_kwargs") or {}) - mm_kwargs["text_kwargs"] = dict(mm_kwargs.get("text_kwargs") or {}) - if Version(TRANSFORMERS_VERSION) < Version("4.58.0"): - # Extract audio_sample_rate before restructuring - audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None) - - # move truncation to audio_kwargs level to avoid conflict - # with tok_kwargs - mm_kwargs["audio_kwargs"].setdefault("truncation", mm_kwargs.pop("truncation", False)) - mm_kwargs["text_kwargs"].setdefault("truncation", tok_kwargs.pop("truncation", False)) - - # Validate and conditionally pass audio_sample_rate - # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's - # audio loader already resamples audio to the target rate. - # Only pass the value if it matches to avoid unexpected behavior. - if audio_sample_rate is not None: - expected_sr = feature_extractor.sampling_rate - if audio_sample_rate != expected_sr: - logger.warning( - "[%s] audio_sample_rate mismatch: user provided %dHz " - "but model expects %dHz. Ignoring user value. " - "vLLM's audio loader already resampled to %dHz.", - self.__class__.__name__, - audio_sample_rate, - expected_sr, - expected_sr, - ) - else: - # Sample rate matches, safe to pass - mm_kwargs["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate - - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - if ( - "audio_feature_lengths" in hf_inputs - and "feature_attention_mask" in hf_inputs - and (audios := mm_data.get("audio", [])) - ): - audio_num_frames = [] - for _, audio in enumerate(audios): - audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio) - num_frame = ( - (audio_length // hop_length) if audio_length % hop_length == 0 else (audio_length // hop_length - 1) - ) - if mm_kwargs.get("truncation", False): - num_frame = min(num_frame, feature_extractor.n_samples // hop_length) - audio_num_frames.append(num_frame) - hf_inputs["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames] - hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames) - return hf_inputs - - def _maybe_apply_prompt_updates( - self, - mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, - is_update_applied: bool, - ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]: - """ - Qwen3-Omni reimplements this function to handle `use_audio_in_video`. - """ - mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - - video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) - use_audio_in_video = any(video_use_audio_in_video) - - # normal case with `use_audio_in_video=False` - if is_update_applied: - mm_placeholders = self._find_mm_placeholders( - prompt_ids, - mm_prompt_updates, - ) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - else: - if use_audio_in_video and "audio" in mm_prompt_updates: - filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - filtered_updates, - ) - mm_placeholders = self._derive_audio_from_video_placeholders( - mm_placeholders, - mm_prompt_updates, - video_use_audio_in_video, - ) - else: - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, - ) - - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - - return prompt_ids, mm_placeholders + # Preserve Omni's per-video overrides while inheriting upstream Qwen3's + # audio processing through the shared, cooperative Qwen2.5 base. def get_updates_use_audio_in_video( self, @@ -978,17 +841,17 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_qwen2_audio, ), PromptReplacement( modality="image", - target=image_token, + target=[image_token_id], replacement=partial(get_replacement_qwen2_vision, modality="image"), ), PromptReplacement( modality="video", - target=video_token, + target=[video_token_id], replacement=get_replacement_qwen2_video, ), ] @@ -1114,7 +977,10 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix audio_input: Qwen2_5OmniAudioFeatureInputs, ) -> tuple[torch.Tensor, ...]: input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] + # vLLM 0.29 marks audio_feature_lengths keep_on_cpu, and the audio tower + # derives device placement from feature_lens, so move it explicitly. + # Mirrors upstream Qwen3OmniMoeThinker._process_audio_input. + audio_feature_lengths = audio_input["audio_feature_lengths"].to(input_features.device, non_blocking=True) audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) @@ -1540,11 +1406,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"talker.": None, "code2wav.": None}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 87df87889..f1692f73d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -11,7 +14,7 @@ from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp @@ -552,10 +555,9 @@ class Qwen3TTSCode2Wav(nn.Module): subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - loaded = AutoWeightsLoader( - self, - skip_prefixes=["encoder."], - ).load_weights(subfolder_weights) + loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"encoder.": None}) + ) device = self.vllm_config.device_config.device self.decoder.to(device=device, dtype=torch.float32) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index e65334ff5..dfb24516a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -1100,10 +1103,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - enc_loaded = AutoWeightsLoader( - self, - skip_prefixes=["decoder."], - ).load_weights(subfolder_weights) + enc_loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"decoder.": None}) + ) loaded |= enc_loaded # AutoWeightsLoader only loads parameters; the encoder's VQ 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 b00e24870..6dca8b358 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 @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding @@ -19,6 +22,7 @@ import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from transformers import AutoTokenizer +from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ProcessorMixin from vllm.config import VllmConfig from vllm.inputs import MultiModalDataDict @@ -503,6 +507,10 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) audio_token = getattr(processor, "audio_token", "") + audio_token_ids = self.info.get_tokenizer().encode( + audio_token, + add_special_tokens=False, + ) audio_items = mm_items.get("audio", []) num_audio_items = len(audio_items) if audio_items else 0 @@ -548,27 +556,22 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing return [ PromptReplacement( modality="audio", - target=audio_token, + target=audio_token_ids, replacement=get_replacement_audio, ) ] - def _call_hf_processor( + def _preprocess_hf_mm_data( self, - prompt: str, mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object] | None = None, - tok_kwargs: Mapping[str, object] | None = None, - **_: object, - ): - """Call HF processor and post-process outputs""" + 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", []) - - mm_kwargs = mm_kwargs or {} - tok_kwargs = tok_kwargs or {} - - # CRITICAL: Ensure audios is ALWAYS a list to prevent string iteration if audios: if isinstance(audios, str): mm_data["audio"] = [audios] @@ -576,20 +579,24 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing mm_data["audio"] = audios else: mm_data["audio"] = [audios] + return mm_data, hf_processor_mm_kwargs - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, + def _postprocess_hf_mm_data( + self, + mm_data: Mapping[str, object], + hf_processor_mm_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + processed_data = super()._postprocess_hf_mm_data( + mm_data, + hf_processor_mm_kwargs, + processed_data, ) - - if "audio_mels" not in hf_inputs: - hf_inputs["audio_mels"] = torch.empty((0, 128, 0)) - if "audio_lens" not in hf_inputs: - hf_inputs["audio_lens"] = torch.tensor([], dtype=torch.int32) - - return hf_inputs + if "audio_mels" not in processed_data: + processed_data["audio_mels"] = torch.empty((0, 128, 0)) + if "audio_lens" not in processed_data: + processed_data["audio_lens"] = torch.tensor([], dtype=torch.int32) + return processed_data # ============================================================================ 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 4f63bbec1..211159afa 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 @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math from collections.abc import Iterable, Mapping, Sequence @@ -48,11 +48,9 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, Mult from vllm.multimodal.processing import BaseDummyInputsBuilder, BaseMultiModalProcessor from vllm.multimodal.processing.processor import ( BaseProcessingInfo, - MultiModalProcessingInfo, ProcessorInputs, PromptReplacement, PromptUpdate, - TimingContext, ) from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config @@ -871,7 +869,7 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing return [ PromptReplacement( modality="audio", - target="", # Never match the prompt (see below note) + target=[audio_id], replacement=get_replacement, ), ] @@ -899,15 +897,24 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing ) return mm_processed_data - def _cached_apply_hf_processor( + def _maybe_apply_prompt_updates( self, - inputs: ProcessorInputs, - timing_ctx: TimingContext, - ) -> tuple[list[int], MultiModalProcessingInfo, bool]: - prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(inputs, timing_ctx) - - # NOTE: The tokens are already inserted by the chat template - return prompt_ids, mm_info, True + mm_items: MultiModalDataItems, + prompt_ids: list[int], + mm_kwargs: MultiModalKwargsItems, + mm_prompt_updates, + ): + # 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() + 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( + prompt_ids, + mm_prompt_updates, + ) + self._validate_mm_placeholders(mm_placeholders, mm_item_counts) + return prompt_ids, mm_placeholders @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm_omni/platforms/__init__.py b/vllm_omni/platforms/__init__.py index 942eb7c24..d5ee64353 100644 --- a/vllm_omni/platforms/__init__.py +++ b/vllm_omni/platforms/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import logging import traceback @@ -142,7 +142,11 @@ def resolve_current_omni_platform_cls_qualname() -> str: if platform_cls_qualname is not None: activated_plugins.append(name) except Exception: - pass + logger.debug( + "OmniPlatform plugin %s failed during detection.", + name, + exc_info=True, + ) activated_builtin_plugins = list(set(activated_plugins) & set(builtin_omni_platform_plugins.keys())) activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys())) diff --git a/vllm_omni/reasoning/step_audio_reasoning_parser.py b/vllm_omni/reasoning/step_audio_reasoning_parser.py index 81ef63496..ddf5a434f 100644 --- a/vllm_omni/reasoning/step_audio_reasoning_parser.py +++ b/vllm_omni/reasoning/step_audio_reasoning_parser.py @@ -5,7 +5,7 @@ from collections.abc import Iterable, Sequence from itertools import islice from typing import TYPE_CHECKING -from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.logger import init_logger from vllm.reasoning import ReasoningParser from vllm.tokenizers import TokenizerLike diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 811f86f17..6625effe8 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -2128,7 +2128,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex elif self.valid_sampled_token_count_event is not None: assert spec_decode_common_attn_metadata is not None next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded( - self.optimistic_seq_lens_cpu, sampled_token_ids, self.requests, self.input_batch, diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index 845881138..39ecae469 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/vllm_omni/worker/sampling_utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Sampling-state guards shared by the GPU and NPU AR model runners.""" import torch @@ -48,7 +48,7 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in if not min_toks: continue needs_rebuild = False - for _, _, stop_tok_ids in min_toks.values(): + for _, _, stop_tok_ids, _ in min_toks.values(): oob = [tok for tok in stop_tok_ids if tok >= logits_vocab] if not oob: continue @@ -63,10 +63,19 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in if needs_rebuild: reqs: list[int] = [] tok_ids: list[int] = [] - for index, (_, _, stop_tok_ids) in min_toks.items(): + restore_reqs: list[int] = [] + restore_tok_ids: list[int] = [] + for index, (_, _, stop_tok_ids, uses_structured_output) in min_toks.items(): reqs.extend([index] * len(stop_tok_ids)) tok_ids.extend(stop_tok_ids) + if uses_structured_output: + restore_reqs.extend([index] * len(stop_tok_ids)) + restore_tok_ids.extend(stop_tok_ids) proc.logits_slice = ( proc._device_tensor(reqs, torch.int32), proc._device_tensor(tok_ids, torch.int32), ) + proc.restore_logits_slice = ( + proc._device_tensor(restore_reqs, torch.int32), + proc._device_tensor(restore_tok_ids, torch.int32), + )