[Rebase] Rebase to vllm 0.26.0 (#5443)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: tzhouam <tzhouam@gmail.com>

.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

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

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

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}}
 

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
 

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"
 

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
 ```
 

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]

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
 

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

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.
 

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,
     )
 

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 == {}

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():

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)
 

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"]

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}"
     )
 

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"
 

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]

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:

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(

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 (

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"})

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
 

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])

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)
 

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,

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"],
             ),
         ]
 

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 <boi>
+                    token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise using default <boi>
                     extra_token_pos["boi"].append(token_count)
                     token_count += 1
                 token_count = self._add_image_meta_info_token(

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

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):

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:

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)
 
 

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

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

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.

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."""

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

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

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

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