From cbab1d4d8f814e48ee691b69ac4cee9d0419f218 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Mon, 27 Jul 2026 23:35:49 +0800 Subject: [PATCH] [Rebase] Rebase to vllm 0.26.0 (#5443) Signed-off-by: tzhouam Signed-off-by: tzhouam Co-authored-by: Claude Fable 5 Co-authored-by: tzhouam --- .buildkite/cuda/test-merge.yml | 8 ++- docker/Dockerfile.ci | 4 +- docker/Dockerfile.cuda | 2 +- docker/Dockerfile.rocm | 4 +- docker/Dockerfile.xpu | 2 +- docs/getting_started/installation/gpu.md | 2 +- .../installation/gpu/cuda.inc.md | 16 +++--- .../installation/gpu/rocm.inc.md | 14 ++--- docs/getting_started/quickstart.md | 6 +-- tests/attention/test_fish_kvcache_attn.py | 53 ++++++++++--------- .../sched/test_omni_ar_scheduler_streaming.py | 20 +++++++ ...st_omni_scheduler_finish_requests_purge.py | 17 +++--- ...mni_scheduler_input_coordinator_cleanup.py | 10 ++-- .../test_boogu_image_transformer.py | 45 +++++++++++----- .../test_chunk_transfer_adapter.py | 12 +++-- .../omni_connectors/test_kv_flow.py | 25 +++++++++ .../test_qwen3_omni_realtime_websocket.py | 12 ++++- tests/entrypoints/test_pd_disaggregation.py | 27 ++-------- .../qwen2_5_omni/test_qwen2_5_omni_embed.py | 47 +++++++++++----- .../qwen3_omni/test_qwen3_omni_moe_lora.py | 1 + tests/worker/test_gpu_ar_model_runner.py | 3 ++ .../test_gpu_generation_model_runner.py | 2 +- tests/worker/test_omni_connector_mixin.py | 2 +- vllm_omni/__init__.py | 10 +++- vllm_omni/attention/fish_kvcache_backend.py | 10 ++-- vllm_omni/core/sched/omni_ar_scheduler.py | 35 +++++++----- .../core/sched/omni_generation_scheduler.py | 29 ++++++---- .../models/cosmos3/pipeline_cosmos3.py | 3 +- .../hunyuan_image3_tokenizer.py | 2 +- .../omni_connectors/kv_transfer_manager.py | 2 +- .../omni_connectors/utils/kv_utils.py | 26 ++++++++- .../ar_diffusion/kv_cache/manager.py | 5 ++ .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 14 +---- .../models/moss_tts/moss_codec_cudagraph.py | 10 ++-- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 2 - .../qwen3_omni/qwen3_omni_moe_thinker.py | 10 +++- vllm_omni/platforms/cuda/platform.py | 10 ++++ vllm_omni/worker/gpu_ar_model_runner.py | 5 +- .../worker/gpu_generation_model_runner.py | 5 +- vllm_omni/worker/gpu_model_runner.py | 3 +- .../omni_connector_model_runner_mixin.py | 2 +- 41 files changed, 334 insertions(+), 183 deletions(-) diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 22da9617b..96ea162b4 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -49,13 +49,17 @@ steps: depends_on: upload-merge-pipeline steps: - label: "Entrypoint Test with H100" - timeout_in_minutes: 60 + # TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per + # run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds + # 2915 vs 2916), which pushed this suite past 60 min. Revert once the + # 0.26 startup regression is understood and fixed. + timeout_in_minutes: 90 commands: - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" mirror_hardwares: h100_2 - label: "Entrypoint Test with L4" - timeout_in_minutes: 30 + timeout_in_minutes: 60 commands: - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model" mirror_hardwares: l4_1 diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index caae57abf..02a7be452 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,7 +1,7 @@ -# Target vllm version: v0.25.0 (releases/v0.25.0 branch). +# Target vllm version: v0.26.0 (releases/v0.26.0 branch). # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.25.0 +ARG VLLM_BASE_TAG=v0.26.0 # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 4e9075f2a..f2e80eaf3 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 9be3e5fb5..6b5a12309 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. @@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base # we should swap over to use stable release ASAP. # We should avoid relying on custom commits. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index b70dd274e..5d298040c 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" -ARG VLLM_VERSION=v0.25.0 +ARG VLLM_VERSION=v0.26.0 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm WORKDIR /workspace/vllm diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index e21a11cee..59d2aa482 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels -Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). +Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. 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 d4808139e..9d568bde3 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.25.0 --torch-backend=auto +uv pip install vllm==0.26.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library ```bash -uv pip install vllm==0.25.0 --torch-backend=auto +uv pip install vllm==0.26.0 --torch-backend=auto ``` -The 0.25.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.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.25.0 +git checkout v0.26.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.25.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.25.0 \ + vllm/vllm-omni:v0.26.0 \ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 4a64e8723..c44407299 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 +uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library ```bash -uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 +uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 ``` -The pre-built 0.25.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.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.25.0 +git checkout v0.26.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -130,7 +130,7 @@ docker run --rm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ - vllm/vllm-omni-rocm:v0.25.0 \ + vllm/vllm-omni-rocm:v0.26.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -149,7 +149,7 @@ docker run --rm -it \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ --entrypoint bash \ - vllm/vllm-omni-rocm:v0.25.0 + vllm/vllm-omni-rocm:v0.26.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 914f1f7b4..15ceb395e 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.25.0 --torch-backend=auto +uv pip install vllm==0.26.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 +uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.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.25.0` and vLLM Omni `0.25.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.26.0` and vLLM Omni `0.26.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 c82fd1858..d38380e13 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.py @@ -100,8 +100,9 @@ def _decode_kv_cache( head_size: int = 128, dtype: torch.dtype = torch.float16, ) -> torch.Tensor: - # vLLM >=0.23.0 KV cache layout: key/value live on dim 1. - return torch.zeros((num_blocks, 2, block_size, num_kv_heads, head_size), dtype=dtype) + # vLLM >=0.23.0 KV cache layout: (num_blocks, num_kv_heads, block_size, + # 2*head_size) with K and V packed in the last dimension. + return torch.zeros((num_blocks, num_kv_heads, block_size, 2 * head_size), dtype=dtype) def test_fish_kvcache_enabled_by_default(monkeypatch): @@ -221,10 +222,11 @@ def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size): def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): - # vLLM >=0.23.0 lays the KV cache out as (num_blocks, 2, block_size, - # num_kv_heads, head_size). The backend must unbind on dim 1 so each of - # key/value comes back as the 4-D (num_blocks, block_size, num_kv_heads, - # head_size) tensor the decode kernel expects. + # vLLM >=0.23.0 lays out the KV cache as (num_blocks, num_kv_heads, + # block_size, 2*head_size) with K and V packed in the last dimension. + # The backend must split on the content dim so each of key/value comes + # back as the 4-D (num_blocks, block_size, num_kv_heads, head_size) + # tensor the decode kernel expects. monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None) @@ -232,7 +234,8 @@ def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128 cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) - assert cache_shape == (num_blocks, 2, block_size, num_kv_heads, head_size) + # Upstream now returns a 4-D packed shape: (B, H, N, 2*D). + assert cache_shape == (num_blocks, num_kv_heads, block_size, 2 * head_size) captured = {} @@ -262,20 +265,21 @@ def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkeypatch): # End-to-end with the real vLLM cache layout and the real (unmocked) guard. # - # The point of the unbind(1) fix: unbind(0) raised "too many values to - # unpack" on the 5-D vLLM 0.23.0 layout, so the request crashed outright. - # unbind(1) makes the request *work* again -- it unpacks cleanly and the - # backend gracefully falls back to the original attention with correct - # output. This test guards that fix and must NOT depend on the fast path - # firing. + # The point of the transpose+split fix: unbind(1) on the old 5-D layout + # raised "too many values to unpack" on the vLLM 0.23.0 layout, and the + # upstream now packs K/V into the last dimension (4-D shape). The + # transpose(1, 2) + split correctly unpacks into separate key/value tensors + # and the backend gracefully falls back to the original attention with + # correct output. This test guards that fix and must NOT depend on the fast + # path firing. # - # It also documents a known limitation (tracked separately): unbind(1) on - # the interleaved (num_blocks, 2, ...) layout yields *non-contiguous* - # key/value views, which the guard rejects on the is_contiguous() check. - # So the dimensions are compatible (4-D, block_size=16, head_size=128) -- - # the only thing keeping the fast path from engaging is contiguity, not - # shape. Making the fast path actually fire would require a .contiguous() - # copy and is out of scope here. + # It also documents a known limitation (tracked separately): + # transpose(1, 2) on the interleaved (B, H, N, 2*D) layout yields + # *non-contiguous* key/value views, which the guard rejects on the + # is_contiguous() check. So the dimensions are compatible (4-D, + # block_size=16, head_size=128) -- the only thing keeping the fast path + # from engaging is contiguity, not shape. Making the fast path actually + # fire would require a .contiguous() copy and is out of scope here. fish_kvcache_backend.reset_fish_kvcache_attn_stats() monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) @@ -300,13 +304,14 @@ def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkey output = torch.zeros_like(query) kv_cache = torch.zeros(cache_shape, dtype=torch.float16) - # Root cause record: unbind(1) gives the right 4-D shape but a non-contiguous view. - key_cache, value_cache = kv_cache.unbind(1) + # Root cause record: transpose(1,2) + split gives the right 4-D shape but + # a non-contiguous view. + key_cache, value_cache = kv_cache.transpose(1, 2).split(head_size, dim=-1) assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size) assert not key_cache.is_contiguous() - # unbind(1) no longer crashes (unbind(0) would have raised here) and the - # request completes via the correct fallback path. + # transpose+split no longer crashes (unbind(1) would have raised here) and + # the request completes via the correct fallback path. result = model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output) assert result is output diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 06037d4cf..4ae4f7684 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -73,6 +73,9 @@ def _run_resumable_segment_stop( sched._update_request_with_output.side_effect = stop_request sched._handle_stopped_request.return_value = session_finished + # vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured + # MagicMock iterates empty and fails to unpack at the call site. + sched._free_request.return_value = (None, None) sched.chunk_transfer_adapter = None sched.running = [session] sched.waiting_for_transfer_free = set() @@ -137,6 +140,23 @@ def test_resumable_session_terminal_is_not_marked_as_segment_boundary() -> None: assert output.is_segment_finished is False +def test_update_from_output_settles_in_flight_tokens() -> None: + """vLLM 0.26: schedule() increments num_in_flight_tokens per scheduled + token; update_from_output must decrement it symmetrically. If the + decrement is dropped the counter grows monotonically and both readers + (allocate_slots, _connector_finished) clamp + max(0, num_computed_tokens - num_in_flight_tokens) to zero forever, + silently freezing sliding-window block freeing. + """ + session = _make_request() + session.status = RequestStatus.RUNNING + session.num_in_flight_tokens = 1 # as left by schedule() for this step + + _run_resumable_segment_stop(session) + + assert session.num_in_flight_tokens == 0 + + def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None: """A decode step that neither stops nor carries an inter-stage payload. diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py index 9b564d7d5..fbfeb384f 100644 --- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py +++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py @@ -76,8 +76,9 @@ def test_finish_requests_purges_finished_request_from_running( ) def fake_finish_requests(self, request_ids, finished_status): - # super() returning the abort tuple but leaving self.running untouched - return [(finished.request_id, 0)] + # vLLM 0.26 super() returns the finished Request objects but here + # leaves self.running untouched to simulate the missed removal. + return [finished] monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) @@ -88,7 +89,7 @@ def test_finish_requests_purges_finished_request_from_running( ) assert scheduler.running == [] - assert result == [(finished.request_id, 0)] + assert result == [finished] @pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) @@ -111,7 +112,7 @@ def test_finish_requests_purges_untracked_request_from_running( ) def fake_finish_requests(self, request_ids, finished_status): - return [(untracked.request_id, 0)] + return [untracked] monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) @@ -122,7 +123,7 @@ def test_finish_requests_purges_untracked_request_from_running( ) assert scheduler.running == [] - assert result == [(untracked.request_id, 0)] + assert result == [untracked] @pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) @@ -136,7 +137,7 @@ def test_finish_requests_leaves_healthy_running_intact( is conservative and won't aggressively drop healthy entries. """ alive = _StubRequest("req-alive", RequestStatus.RUNNING) - leaving_id = "req-leaving" + leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED) scheduler = _make_scheduler( scheduler_cls, requests={alive.request_id: alive}, @@ -145,13 +146,13 @@ def test_finish_requests_leaves_healthy_running_intact( ) def fake_finish_requests(self, request_ids, finished_status): - return [(leaving_id, 0)] + return [leaving] monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) scheduler_cls.finish_requests( scheduler, - [leaving_id], + [leaving.request_id], RequestStatus.FINISHED_STOPPED, ) diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index fcac8172b..cb0f50a5d 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -53,7 +53,11 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids( def fake_finish_requests(self, request_ids, finished_status): assert request_ids == ["req-a", "req-b"] assert finished_status == RequestStatus.FINISHED_STOPPED - return [("req-a", 0), ("req-b", 1)] + # vLLM 0.26 returns the finished Request objects, not (req_id, client_index) tuples. + return [ + SimpleNamespace(request_id="req-a", client_index=0), + SimpleNamespace(request_id="req-b", client_index=1), + ] monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) @@ -63,7 +67,7 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids( RequestStatus.FINISHED_STOPPED, ) - assert result == [("req-a", 0), ("req-b", 1)] + assert [(r.request_id, r.client_index) for r in result] == [("req-a", 0), ("req-b", 1)] assert coordinator.freed == ["req-a", "req-b"] @@ -82,7 +86,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None: result = OmniARScheduler._free_request(scheduler, DummyRequest()) - assert result is None + assert result == (None, None) assert coordinator.freed == ["req-free"] assert scheduler._omits_kv_transfer_cache == {} assert scheduler._new_prompt_len_snapshot == {} 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 5551ba191..08ba4af09 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging import os from types import SimpleNamespace @@ -338,27 +339,43 @@ def test_transformer_load_weights_round_trip(): assert torch.allclose(reloaded[native_name], expected) -def test_transformer_load_weights_warns_for_unexpected_and_unloaded(caplog): +def test_transformer_load_weights_warns_for_unexpected_and_unloaded(): from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( BooguImageTransformer2DModel, ) - model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) - native_params = dict(model.named_parameters()) - loaded_name = next(iter(native_params)) - checkpoint_name = _native_to_checkpoint_name(loaded_name) + # vllm_omni's logger hierarchy hangs off vLLM's `vllm` logger, which sets + # propagate=False, so records never reach the root logger that pytest's + # caplog listens on. Capture at the emitting logger instead. + messages: list[str] = [] - loaded = model.load_weights( - [ - (checkpoint_name, torch.randn_like(native_params[loaded_name])), - ("unexpected.weight", torch.ones(1)), - ] - ) + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + messages.append(record.getMessage()) + log = logging.getLogger("vllm_omni.diffusion.models.boogu_image.boogu_image_transformer") + handler = _Capture(level=logging.WARNING) + log.addHandler(handler) + try: + model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) + native_params = dict(model.named_parameters()) + loaded_name = next(iter(native_params)) + checkpoint_name = _native_to_checkpoint_name(loaded_name) + + loaded = model.load_weights( + [ + (checkpoint_name, torch.randn_like(native_params[loaded_name])), + ("unexpected.weight", torch.ones(1)), + ] + ) + finally: + log.removeHandler(handler) + + text = "\n".join(messages) assert loaded == {loaded_name} - assert "Skipping unexpected checkpoint weight unexpected.weight" in caplog.text - assert "Model parameters not loaded from checkpoint" in caplog.text - assert next(name for name in native_params if name != loaded_name) in caplog.text + assert "Skipping unexpected checkpoint weight unexpected.weight" in text + assert "Model parameters not loaded from checkpoint" in text + assert next(name for name in native_params if name != loaded_name) in text def test_transformer_forward_t2i_shape(): diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 1485ace60..1a045f306 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -780,6 +780,10 @@ def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_ada class _HashableRequest(SimpleNamespace): """SimpleNamespace that can be added to a set (needed by scheduler internals).""" + # vLLM 0.26: update_from_output settles this counter for every scheduled + # request; real Requests initialise it to 0 (vllm/v1/request.py). + num_in_flight_tokens = 0 + def __hash__(self): return hash(self.request_id) @@ -833,7 +837,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc scheduler.requests = {"req-s1": request} scheduler._handle_stopped_request = mocker.MagicMock(return_value=True) - scheduler._free_request = mocker.MagicMock(return_value=None) + scheduler._free_request = mocker.MagicMock(return_value=(None, None)) scheduler._get_routed_experts = mocker.MagicMock(return_value=None) scheduler.running = [request] scheduler.waiting = mocker.MagicMock() @@ -918,7 +922,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True)) scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False) scheduler._handle_stopped_request = mocker.MagicMock(return_value=True) - scheduler._free_request = mocker.MagicMock(return_value=None) + scheduler._free_request = mocker.MagicMock(return_value=(None, None)) scheduler._get_routed_experts = mocker.MagicMock(return_value=None) scheduler.running = [request] scheduler.waiting = mocker.MagicMock() @@ -1047,7 +1051,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs): scheduler._pending_finish_reqs = list(pending_finish_reqs) scheduler._handle_stopped_request = mocker.MagicMock(return_value=True) - scheduler._free_request = mocker.MagicMock(return_value=None) + scheduler._free_request = mocker.MagicMock(return_value=(None, None)) scheduler._get_routed_experts = mocker.MagicMock(return_value=None) scheduler.running = list(running) scheduler.waiting = mocker.MagicMock() @@ -1103,7 +1107,7 @@ def test_deferred_finish_emits_finished_output(mocker: MockerFixture): running=[request], pending_finish_reqs=[request], ) - scheduler._free_request.return_value = {"mock": "kv_params"} + scheduler._free_request.return_value = ({"mock": "kv_params"}, None) result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out) diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 83206fefb..b738c095e 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -330,6 +330,31 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants, assert normalized is None +def test_normalize_layer_kv_unpacks_packed_layout(common_constants): + """The vLLM >= 0.23 packed layout (num_blocks, n_heads, block_size, + 2*head_dim) is unpacked into (num_blocks, block_size, n_heads, head_dim) + key/value blocks when block_size is provided.""" + block_size = common_constants["block_size"] + num_heads = common_constants["num_heads"] + head_dim = common_constants["head_dim"] + req_id = common_constants["req_id"] + num_blocks = 10 + + layer_kv = torch.randn(num_blocks, num_heads, block_size, 2 * head_dim) + + # Without block_size the 4-D packed layout is not recognized. + assert normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) is None + + normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0, block_size=block_size) + assert normalized is not None + key_blocks, value_blocks = normalized + assert key_blocks.shape == (num_blocks, block_size, num_heads, head_dim) + assert value_blocks.shape == (num_blocks, block_size, num_heads, head_dim) + reference = layer_kv.transpose(1, 2) + assert torch.equal(key_blocks, reference[..., :head_dim]) + assert torch.equal(value_blocks, reference[..., head_dim:]) + + def test_manager_reception(kv_config, mock_connector, common_constants): """Test reception and injection logic in OmniKVTransferManager.""" num_layers = common_constants["num_layers"] diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 50f8eef91..325daabdd 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -201,6 +201,7 @@ def _assert_realtime_smoke(result: dict) -> None: def _assert_realtime_accuracy( result: dict, whisper_model_size: str = "large-v3", + threshold: float = 0.8, ) -> None: """Assert that whisper transcription of audio output matches model text. @@ -214,6 +215,12 @@ def _assert_realtime_accuracy( variability even though audio generation was correct. large-v3 transcribes these clips reliably, so a failure here now points at the model, not the ASR grader. + threshold: Minimum cosine similarity (with length penalty) required to + pass. Default 0.8. Do not lower per-callsite without data: + at 0.35 the assertion no longer detects real audio + regressions. If a variant genuinely needs a different gate + (e.g. whisper partial transcripts under async_chunk), propose + it in its own PR with measurements. """ final_text = (result["transcription_text"] or "").strip() assert final_text, "Expected non-empty transcription (model text stream)" @@ -223,8 +230,9 @@ def _assert_realtime_accuracy( assert whisper_text, "Whisper returned empty string for synthesized output audio" sim = cosine_similarity_text(whisper_text.lower(), final_text.lower()) - assert sim > 0.8, ( - f"Output audio transcript should match model text (sim={sim:.3f}): " + assert sim > threshold, ( + f"Output audio transcript should match model text (sim={sim:.3f}, " + f"threshold={threshold}): " f"whisper={whisper_text!r}, model_text={final_text!r}" ) diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index 57e484af1..774e6d1a8 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -305,21 +305,11 @@ def mock_get_config(monkeypatch): """Auto-mock get_config and related model loading functions.""" import sys - fake_tokenizer = _ns() - fake_tokenizer.encode = lambda *args, **kwargs: [1, 2, 3] - fake_tokenizer.decode = lambda *args, **kwargs: "test" - - def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): - return fake_tokenizer - - monkeypatch.setattr( - "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs", - _mock_init_tokenizer_from_configs, - raising=False, - ) - tokenizer_module_path = "vllm.transformers_utils.tokenizer" - if tokenizer_module_path in sys.modules: - setattr(sys.modules[tokenizer_module_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) + # TODO(pd-unskip): the tokenizer mock that lived here targeted + # vllm.transformers_utils.tokenizer.init_tokenizer_from_configs, which was + # deleted in vLLM 0.26 (the raising=False setattr made it a silent no-op). + # When the module-level skip is lifted, mock whatever tokenizer seam the + # revived PD entrypoint code actually uses. def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None): if prompt_token_ids is not None: @@ -344,13 +334,6 @@ def mock_get_config(monkeypatch): _mock_length_from_prompt_token_ids_or_embeds, ) - monkeypatch.setattr( - "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False - ) - async_omni_path = "vllm_omni.entrypoints.async_omni" - if async_omni_path in sys.modules: - setattr(sys.modules[async_omni_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) - fake_hf_config = _ns() fake_hf_config.model_type = "qwen2_5_omni" diff --git a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py index d4778b7e2..13899d47d 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py +++ b/tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py @@ -15,6 +15,7 @@ import functools import pytest import torch from pytest_mock import MockerFixture +from vllm.multimodal.utils import set_mm_embedding_modality pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -117,13 +118,30 @@ class TestCheckInterleavedAudioVideo: # --------------------------------------------------------------------------- -def make_mock_model(mocker: MockerFixture, hidden: int = 8): +@functools.lru_cache(maxsize=1) +def _omni_qwen2_5_thinker_cls(): + """vllm_omni's thinker override, loaded lazily like the upstream one.""" + from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerForConditionalGeneration, + ) + + return Qwen2_5OmniThinkerForConditionalGeneration + + +def make_mock_model(mocker: MockerFixture, hidden: int = 8, thinker_cls=None): """ Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration that has enough structure to run embed_input_ids. + + ``thinker_cls`` selects whose ``embed_input_ids`` the mock dispatches to. + It defaults to upstream's. Only the interleaved case passes omni's override: + omni's non-interleaved path falls through to a zero-arg ``super()`` call, + which raises TypeError when ``self`` is a Mock rather than an instance. """ m = _qwen2_5_omni_thinker_mod() Qwen2_5OmniThinkerForConditionalGeneration = m.Qwen2_5OmniThinkerForConditionalGeneration + if thinker_cls is None: + thinker_cls = Qwen2_5OmniThinkerForConditionalGeneration model = mocker.Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration) @@ -155,9 +173,7 @@ def make_mock_model(mocker: MockerFixture, hidden: int = 8): is_multimodal=is_multimodal, ) - model.embed_input_ids = lambda *a, **kw: m.Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids( # noqa: E501 - model, *a, **kw - ) + model.embed_input_ids = lambda *a, **kw: thinker_cls.embed_input_ids(model, *a, **kw) model._super_embed_input_ids = fake_super_embed @@ -171,11 +187,11 @@ def build_mm_embeds(audio_n, image_n, video_n, hidden, audio_val=10.0, image_val """ embs = [] if audio_n: - embs.append(torch.full((audio_n, hidden), audio_val)) + embs.append(set_mm_embedding_modality(torch.full((audio_n, hidden), audio_val), "audio")) if image_n: - embs.append(torch.full((image_n, hidden), image_val)) + embs.append(set_mm_embedding_modality(torch.full((image_n, hidden), image_val), "image")) if video_n: - embs.append(torch.full((video_n, hidden), video_val)) + embs.append(set_mm_embedding_modality(torch.full((video_n, hidden), video_val), "video")) return embs @@ -259,6 +275,11 @@ class TestEmbedInputIds: """ Interleaved (use_audio_in_video): video chunks interleaved with audio. Video embeddings must go to video positions, audio to audio positions. + + Runs against **omni's** thinker override, not upstream's: this is the + only unit-level coverage of omni's call into + ``merge_interleaved_embeddings``, whose upstream signature dropped + ``num_video``/``num_audio`` in vLLM #46213. """ hidden = 8 audio_val, video_val = 10.0, 30.0 @@ -270,11 +291,11 @@ class TestEmbedInputIds: audio_n = sum(audio_chunks) # 6 mm_embeds = [ - torch.full((video_n, hidden), video_val), - torch.full((audio_n, hidden), audio_val), + set_mm_embedding_modality(torch.full((video_n, hidden), video_val), "video"), + set_mm_embedding_modality(torch.full((audio_n, hidden), audio_val), "audio"), ] - model, _ = make_mock_model(mocker, hidden) + model, _ = make_mock_model(mocker, hidden, thinker_cls=_omni_qwen2_5_thinker_cls()) result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal) video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0] @@ -308,8 +329,8 @@ class TestMergeInterleavedEmbeddings: inputs_embeds = torch.zeros(len(input_ids), hidden) mm_embeds = [ - torch.full((num_video, hidden), 30.0), - torch.full((num_audio, hidden), 10.0), + set_mm_embedding_modality(torch.full((num_video, hidden), 30.0), "video"), + set_mm_embedding_modality(torch.full((num_audio, hidden), 10.0), "audio"), ] result = m.merge_interleaved_embeddings( @@ -318,8 +339,6 @@ class TestMergeInterleavedEmbeddings: is_video, is_audio, is_multimodal, - num_video, - num_audio, ) video_pos = is_video.nonzero(as_tuple=True)[0] 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 e5fa38a6b..75d81eed5 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 @@ -84,6 +84,7 @@ class _CpuFusedMoE3DWithLoRA(FusedMoE3DWithLoRA): self.tp_size = 1 self.tp_rank = 0 self._w13_slices = 1 + self.enable_moe_shared_loras = False def _save_peft_adapter(tmp_path: Path) -> Path: diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index a1f885c15..050063978 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -69,6 +69,9 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], + # vLLM 0.26: one entry per KV cache group; max_model_len=8 with + # block_sizes=[1] means one group of 8 blocks. + max_num_blocks_per_req=[8], ) input_batch.add_request( CachedRequestState( diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index e64eecb57..a6df4732f 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -4,7 +4,7 @@ from types import SimpleNamespace import pytest import torch from vllm.config import CacheConfig -from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT import vllm_omni.worker.gpu_generation_model_runner as gen_runner_module from vllm_omni.worker.gpu_generation_model_runner import ( diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 87cf4ba73..cf4e09cb9 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -1245,7 +1245,7 @@ class TestRankAwareKVRouting(unittest.TestCase): class TestAttachOmniConnectorOutput(unittest.TestCase): def test_wraps_empty_model_runner_output_when_signals_exist(self): - from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT + from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT host = MixinHost() host.get_omni_connector_output = lambda: OmniConnectorOutput(chunk_ready_req_ids={"req-1"}) diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index f3c9d9afd..c2cbe9d5c 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -27,8 +27,14 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency patch = None # type: ignore # Register custom configs (AutoConfig, AutoTokenizer) as early as possible. -from vllm_omni.transformers_utils import configs as _configs # noqa: F401, E402 -from vllm_omni.transformers_utils import parsers as _parsers # noqa: F401, E402 +try: + from vllm_omni.transformers_utils import configs as _configs # noqa: F401, E402 + from vllm_omni.transformers_utils import parsers as _parsers # noqa: F401, E402 +except ModuleNotFoundError as exc: # pragma: no cover - optional dependency + if exc.name != "vllm": + raise + # Allow importing vllm_omni without vllm (e.g., documentation builds) + pass from .config import OmniModelConfig diff --git a/vllm_omni/attention/fish_kvcache_backend.py b/vllm_omni/attention/fish_kvcache_backend.py index 0178f722e..4ca6d467b 100644 --- a/vllm_omni/attention/fish_kvcache_backend.py +++ b/vllm_omni/attention/fish_kvcache_backend.py @@ -234,11 +234,11 @@ def _forward_with_fish_kvcache( if attn_metadata is not None and not attn_metadata.use_cascade: num_actual_tokens = attn_metadata.num_actual_tokens # vLLM >=0.23.0 lays out the KV cache as - # (num_blocks, 2, block_size, num_kv_heads, head_size), so key/value - # live on dim 1 (matches FlashAttentionImpl's kv_cache.unbind(1)). - # The pre-0.23 layout had them on dim 0; unbind(0) here would split - # along num_blocks and raise "too many values to unpack". - key_cache, value_cache = kv_cache.unbind(1) + # (num_blocks, num_kv_heads, block_size, 2*head_size) with K and V + # packed in the last dimension. Match FlashAttentionImpl's own + # unpacking: transpose(1, 2) then split on the content dim. + head_size = kv_cache.shape[-1] // 2 + key_cache, value_cache = kv_cache.transpose(1, 2).split(head_size, dim=-1) q = query[:num_actual_tokens] out = output[:num_actual_tokens] active_batch = int(q.shape[0]) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 3892bd389..22917197d 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections import defaultdict +from collections.abc import Iterable from time import time from typing import Any @@ -341,10 +342,18 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is not None: + # vLLM 0.26: settle the in-flight tokens counted in schedule(). + # Must happen before the skips below — failed-KV-load and + # already-finished requests were incremented too, and the two + # readers (allocate_slots, _connector_finished) clamp with + # max(0, computed - in_flight), so a leaked counter silently + # freezes sliding-window block freeing. + request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # Skip requests that were recovered from KV load failure continue - request = self.requests.get(req_id) if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., @@ -451,7 +460,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): request.spec_token_ids = [] request._output_token_ids.clear() if finished: - kv_transfer_params = self._free_request(request) + kv_transfer_params, _ = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: @@ -642,7 +651,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return engine_core_outputs - def finish_requests(self, request_ids: Any, finished_status: RequestStatus) -> list[tuple[str, int]]: + def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]: """Handles the finish signal from outside the scheduler. For example, the API server can abort a request when the client @@ -651,8 +660,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): If request_ids is None, all requests will be finished. Returns: - Tuple of (req_id, client_index) for requests that were aborted. Will not - include any that were already finished. + The Request objects that were aborted. Will not include any that + were already finished. """ # TODO(yrr): chunk transfer adapter & input_coordinator unified to one if self.chunk_transfer_adapter: @@ -688,8 +697,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is not None: - for request_id, _ in finished: - self._free_input_coordinator_request(request_id) + for request in finished: + self._free_input_coordinator_request(request.request_id) return finished def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: @@ -740,7 +749,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer - def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred """Mark a request as finished and free its resources.""" assert request.is_finished() @@ -779,10 +790,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.") self.waiting_for_transfer_free.add(request_id) kv_xfer_params = None - return kv_xfer_params + return kv_xfer_params, None elif request_id in self.waiting_for_transfer_free: # Blocks held until KV extraction completes in a future step. - return None + return None, None else: logger.debug( f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). " @@ -819,14 +830,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if isinstance(add_info, dict): add_info.update(kv_xfer_params) - return kv_xfer_params + return kv_xfer_params, None # 3. Standard Freeing delay_free_blocks |= connector_delay_free_blocks if not delay_free_blocks: self._free_blocks(request) - return kv_xfer_params + return kv_xfer_params, None finally: self._free_input_coordinator_request(request_id) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 534019513..06082b0dc 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -2,6 +2,7 @@ from __future__ import annotations import time from collections import defaultdict +from collections.abc import Iterable from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat @@ -348,7 +349,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): return self._wrap_omni_scheduler_output(scheduler_output) - def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]: + def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]: """Handles the finish signal from outside the scheduler. For example, the API server can abort a request when the client @@ -357,8 +358,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): If request_ids is None, all requests will be finished. Returns: - Tuple of (req_id, client_index) for requests that were aborted. Will not - include any that were already finished. + The Request objects that were aborted. Will not include any that + were already finished. """ if self.chunk_transfer_adapter: @@ -382,11 +383,13 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._purge_finished_from_running() if self.input_coordinator is not None: - for request_id, _ in finished: - self._free_input_coordinator_request(request_id) + for request in finished: + self._free_input_coordinator_request(request.request_id) return finished - def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: if self.input_coordinator is None: return super()._free_request(request, delay_free_blocks) @@ -432,10 +435,18 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is not None: + # vLLM 0.26: settle the in-flight tokens counted in schedule(). + # Must happen before the skips below — failed-KV-load and + # already-finished requests were incremented too, and the two + # readers (allocate_slots, _connector_finished) clamp with + # max(0, computed - in_flight), so a leaked counter silently + # freezes sliding-window block freeing. + request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # Skip requests that were recovered from KV load failure continue - request = self.requests.get(req_id) if request is None or request.is_finished(): # Request may already be finished (e.g., aborted during # execution / pipeline parallelism / async scheduling). @@ -510,7 +521,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.segment_finished_requests.discard(req_id) if finished: - kv_transfer_params = self._free_request(request) + kv_transfer_params, _ = self._free_request(request) if self.chunk_transfer_adapter is not None: self.chunk_transfer_adapter.cleanup( request.request_id, @@ -576,7 +587,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): is_segment_finished = not finished kv_transfer_params = None if finished: - kv_transfer_params = self._free_request(request) + kv_transfer_params, _ = self._free_request(request) if self.chunk_transfer_adapter is not None: self.chunk_transfer_adapter.cleanup( request.request_id, diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 41a3e80fc..fcc2a8b6c 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -956,11 +956,10 @@ class Cosmos3OmniDiffusersPipeline( self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=model_path, - subfolder=None, + subfolder="transformer", revision=None, prefix="transformer.", fall_back_to_pt=True, - allow_patterns_overrides=["transformer/*.safetensors"], ), ] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py index 751bfb21a..e7bb75fbe 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py @@ -384,7 +384,7 @@ class TokenizerWrapper: drop_last_break = True break if source.get("front_boi", use_front_boi_token): - token_seq.append(self.boi_token_id) # Use patched boi for Janus, otherwise useing default + token_seq.append(self.boi_token_id) # Use patched boi for Janus, otherwise using default extra_token_pos["boi"].append(token_count) token_count += 1 token_count = self._add_image_meta_info_token( diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index b75a7089a..efed75241 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -1023,7 +1023,7 @@ class OmniKVTransferManager: value_cache: list[torch.Tensor | None] = [None] * num_layers for layer_idx, layer_kv in enumerate(kv_caches): - kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx) + kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx, block_size=block_size) if kv_pair is None: continue key_blocks, value_blocks = kv_pair diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index fe0df21f2..b7f51460c 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -402,6 +402,7 @@ def normalize_layer_kv( *, req_id: str = "", layer_idx: int = -1, + block_size: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor] | None: """Normalize one layer KV cache to a ``(key_blocks, value_blocks)`` tuple. @@ -424,6 +425,10 @@ def normalize_layer_kv( dim-0 selects key / value. * **Stacked tensor** ``[num_blocks, 2, block_size, n_heads, head_dim]`` – dim-1 selects key / value. + * **Packed tensor** ``[num_blocks, n_heads, block_size, 2 * head_dim]`` – + vLLM >= 0.23 FlashAttention packs K and V into the last (content) + dimension. Only recognized when *block_size* is provided and matches + dim 2, to avoid misreading other 4-D layouts. * **Tuple** ``(key_tensor, value_tensor)`` – returned as-is after validation. @@ -431,13 +436,29 @@ def normalize_layer_kv( layer_kv: The raw KV cache (tensor or tuple) for the layer. req_id: Request ID used only for diagnostic log messages. layer_idx: Layer index used only for diagnostic log messages. + block_size: The paged-attention block size, required to recognize + the packed 4-D layout. Returns: ``(key_blocks, value_blocks)`` if *layer_kv* is valid, ``None`` otherwise. """ if isinstance(layer_kv, torch.Tensor): - if layer_kv.ndim >= 3 and layer_kv.shape[0] == 2: + if ( + layer_kv.ndim == 4 + and block_size is not None + and layer_kv.shape[2] == block_size + and layer_kv.shape[-1] % 2 == 0 + ): + # vLLM >= 0.23 FlashAttention layout: (num_blocks, n_heads, + # block_size, 2 * head_dim). Unpack the same way + # FlashAttentionImpl does: move block_size before heads, then + # split the content dim into the K and V halves. Checked before + # the stacked layouts (which are 5-D in practice) so that a + # 2-KV-head packed tensor is not misread as a dim-1 stack. + head_dim = layer_kv.shape[-1] // 2 + key_blocks, value_blocks = layer_kv.transpose(1, 2).split(head_dim, dim=-1) + elif layer_kv.ndim >= 3 and layer_kv.shape[0] == 2: key_blocks = layer_kv[0] value_blocks = layer_kv[1] elif layer_kv.ndim >= 3 and layer_kv.shape[1] == 2: @@ -446,7 +467,8 @@ def normalize_layer_kv( else: logger.warning( f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: " - f"expected [2, ...] or [..., 2, ...] at dim 0/1, got {tuple(layer_kv.shape)}" + f"expected [2, ...], [..., 2, ...] at dim 0/1, or packed " + f"[num_blocks, n_heads, block_size, 2*head_dim], got {tuple(layer_kv.shape)}" ) return None elif isinstance(layer_kv, tuple): diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index 8c7d4b2ce..9ef3cca0c 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -71,6 +71,11 @@ class ARDiffusionRequestAdapter: self.num_preemptions = 0 # vLLM watermark gate reads this; map the request lifecycle onto it. self.status = RequestStatus.WAITING + # vLLM 0.26 KVCacheManager.allocate_slots reads this when computing + # remove_skipped_blocks. Always 0 here: an AR-Diffusion request only + # advances num_computed_tokens on on_chunk_committed(), so there are + # never optimistically-counted in-flight tokens to subtract. + self.num_in_flight_tokens = 0 @property def num_computed_tokens(self) -> int: 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 6f6841db3..577ce5aa7 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 @@ -105,11 +105,6 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, NestedTensors, ) - -try: - from vllm.multimodal.inputs import ModalityData, MultiModalDataDict -except ImportError: - from vllm.multimodal.parse import ModalityData, MultiModalDataDict from vllm.multimodal.parse import ( AudioItem, AudioProcessorItems, @@ -132,15 +127,10 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors -try: - from vllm.transformers_utils.tokenizer import encode_tokens as _vllm_encode_tokens -except ImportError: - _vllm_encode_tokens = None - +# vllm.transformers_utils.tokenizer no longer exists in upstream vLLM; +# _encode_tokens uses tokenizer.encode() as the only code path. def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: - if _vllm_encode_tokens is not None: - return _vllm_encode_tokens(tokenizer, prompt) return tokenizer.encode(prompt, add_special_tokens=False) diff --git a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py index 87f255072..ed03c018e 100644 --- a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py +++ b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py @@ -1,6 +1,6 @@ """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder. -Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count +Captures MossAudioTokenizerModel._decode for a set of fixed frame-count bucket sizes, then replays the captured graph at inference time to eliminate kernel-launch overhead. Inputs that exceed all captured sizes fall back to eager execution transparently. @@ -24,7 +24,7 @@ logger = init_logger(__name__) class MossTTSCUDAGraphCodecWrapper: - """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame. + """CUDA Graph wrapper for MossAudioTokenizerModel._decode. Graphs are keyed by padded_T (int). On each call the actual T is bucket-matched to the smallest pre-captured size >= T. The static code @@ -99,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper: dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device) dummy_lengths = torch.tensor([size], dtype=torch.long, device=device) with torch.no_grad(): - _ = self.model._decode_frame(dummy_codes, dummy_lengths) + _ = self.model._decode(dummy_codes, dummy_lengths) torch.accelerator.synchronize(device) @@ -132,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper: # Extra eager warmup inside capture to ensure all kernels are compiled. with torch.no_grad(): - _ = self.model._decode_frame(static_codes, static_lengths) + _ = self.model._decode(static_codes, static_lengths) torch.accelerator.synchronize(device) graph = CUDAGraph() with torch.no_grad(): with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): - static_out = self.model._decode_frame(static_codes, static_lengths) + static_out = self.model._decode(static_codes, static_lengths) self.graphs[size] = graph self.static_codes[size] = static_codes 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 827f355a6..2648cacbd 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 @@ -1342,8 +1342,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( is_video, is_audio, is_multimodal, - num_video, - num_audio, ) # Default: standard merge (no interleaving), same as parent class 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 65ec102a1..f7d0553d4 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 @@ -114,6 +114,7 @@ from vllm.multimodal.processing.processor import ( PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.utils import set_mm_embedding_modality from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config @@ -1420,6 +1421,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( visual_dim = embeddings.shape[-1] // (multiscale_len + 1) multi_dim = visual_dim * multiscale_len embeddings_main, embeddings_multiscale = torch.split(embeddings, [visual_dim, multi_dim], dim=-1) + # torch.split returns fresh views, so the `modality` tag the + # encoder gather attached (gpu_model_runner._gather_mm_embeddings) + # does not survive. merge_interleaved_embeddings groups by that + # tag since vLLM #46213, so carry it across the split. + modality = getattr(embeddings, "modality", None) + if modality is not None: + set_mm_embedding_modality(embeddings_main, modality) multimodal_embeddings[index] = embeddings_main multimodal_embeddings_multiscale.append(embeddings_multiscale) if not is_interleaved: @@ -1457,8 +1465,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( is_video, is_audio, is_mm_device, - num_video, - num_audio, ) # Default: standard merge (no interleaving), same as parent class. diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 4f203cfc2..5eb7406a5 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -287,6 +287,16 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): def get_device_name(cls, device_id: int = 0) -> str: return torch.cuda.get_device_name(device_id) + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + device_props = torch.cuda.get_device_properties(device_id) + return device_props.total_memory + + @classmethod + def is_fully_connected(cls, device_ids: list[int]) -> bool: + logger.debug("NVLink detection not available on CudaOmniPlatform; assuming no NVLink.") + return False + @classmethod def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig: """Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops.""" diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index ca7c2a21a..7d9580c91 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -22,9 +22,10 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_ from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer @@ -33,9 +34,7 @@ from vllm.v1.spec_decode.gemma4 import Gemma4Proposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( - EMPTY_MODEL_RUNNER_OUTPUT, AsyncGPUModelRunnerOutput, - IntermediateTensors, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import is_residual_scattered_for_sp diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 2c81643dd..3c00fad38 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -18,9 +18,10 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context +from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer @@ -28,9 +29,7 @@ from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesPropose from vllm.v1.spec_decode.gemma4 import Gemma4Proposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( - EMPTY_MODEL_RUNNER_OUTPUT, AsyncGPUModelRunnerOutput, - IntermediateTensors, PerLayerAttnMetadata, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index cb52a2b67..9fb781dc4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -13,6 +13,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType +from vllm.sequence import IntermediateTensors from vllm.tracing import instrument from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv @@ -26,7 +27,7 @@ from vllm.v1.spec_decode.ngram_proposer_gpu import ( update_scheduler_for_invalid_drafts, ) from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata +from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_omni.core.prefix_cache import OmniTensorPrefixCache diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index b31b04688..363b850d4 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -1575,7 +1575,7 @@ class OmniConnectorModelRunnerMixin: from copy import copy - from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT + from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT wrapped = copy(result if result is not None else EMPTY_MODEL_RUNNER_OUTPUT) wrapped.omni_connector_output = omni_output