[Rebase] Rebase to vllm v0.21.0 (#3530)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Signed-off-by: Dnoob <dxpouo@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: rein yang <ruiruyang2@gmail.com>
Signed-off-by: Nick Cao <ncao@redhat.com>
Signed-off-by: zhumingjue <zhumingjue@huawei.com>
Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>
Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: weizhoublue <weizhoublue@github.com>
Signed-off-by: weizhou.lan@daocloud.io <weizhou.lan@daocloud.io>
Signed-off-by: dengyunyang <584797741@qq.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Dnoob <dxpouo@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: rein yang <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: Nick Cao <ncao@redhat.com>
Co-authored-by: zhumingjue138 <zhumingjue@huawei.com>
Co-authored-by: Ricardo Noriega <rnoriega@redhat.com>
Co-authored-by: lyj-jjj <liuyingjun5@huawei.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: weizhoublue <45163302+weizhoublue@users.noreply.github.com>
Co-authored-by: weizhoublue <weizhoublue@github.com>
Co-authored-by: dengyunyang <584797741@qq.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: NATURE <wzliu@connect.hku.hk>
Co-authored-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Y. Fisher <yukexiong1@huawei.com>
Co-authored-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

tests/diffusion/attention/test_attention_config.py

@@ -394,6 +394,9 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
                 per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
             ),
             parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
         )
 
         with set_current_diffusion_config(od_config):

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
@@ -21,6 +22,11 @@ IMAGE_VIDEO_MODELS = {
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
+_GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
+
+# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py).
+_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible
+
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
@@ -94,10 +100,23 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 @pytest.mark.parametrize("model_name", list(MODELS.keys()))
 def test_cpu_offload_diffusion_model(model_name: str):
+    if model_name in _GATED_MODELS:
+        _skip_if_gated_repo_inaccessible(model_name)
     try:
         offload_peak_memory, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
         no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
+    except ValueError as exc:
+        # omni_snapshot_download wraps GatedRepoError in a ValueError.
+        # If the pre-flight guard above did not catch it (e.g. partial
+        # HF_TOKEN where config.json is accessible but weight shards are
+        # blocked), skip instead of failing.
+        if "Access to model" in str(exc) and "is restricted" in str(exc):
+            pytest.skip(
+                f"Skipping: gated HF repo {model_name!r} inaccessible "
+                f"({exc}). See docs/contributing/ci/hf_credentials.md."
+            )
+        pytest.fail(f"Inference failed: {exc}")
     except Exception:
         pytest.fail("Inference failed")
     print(f"Offload peak memory: {offload_peak_memory} MB")

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -3,6 +3,10 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.diffusion.offloader.test_diffusion_cpu_offload import (
+    _GATED_MODELS,
+    _skip_if_gated_repo_inaccessible,
+)
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -96,6 +100,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     offloader keeps only a single transformer block on GPU at a time, with
     prefetching for compute-memory overlap.
     """
+    if model_name in _GATED_MODELS:
+        _skip_if_gated_repo_inaccessible(model_name)
     try:
         # Run without layerwise offloading (baseline)
         no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
@@ -104,6 +110,17 @@ def test_layerwise_offload_diffusion_model(model_name: str):
         # Run with layerwise offloading (1 layer on device)
         layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
         cleanup_dist_env_and_memory()
+    except ValueError as exc:
+        # omni_snapshot_download wraps GatedRepoError in a ValueError.
+        # If the pre-flight guard above did not catch it (e.g. partial
+        # HF_TOKEN where config.json is accessible but weight shards are
+        # blocked), skip instead of failing.
+        if "Access to model" in str(exc) and "is restricted" in str(exc):
+            pytest.skip(
+                f"Skipping: gated HF repo {model_name!r} inaccessible "
+                f"({exc}). See docs/contributing/ci/hf_credentials.md."
+            )
+        pytest.fail(f"Inference failed: {exc}")
     except Exception:
         pytest.fail("Inference failed")
 

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -540,6 +540,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         kv_connector_output=None,
         cudagraph_stats=None,
         req_id_to_index={"req-s1": 0},
+        routed_experts_dict=None,
     )
 
     OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
@@ -625,6 +626,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         cudagraph_stats=None,
         req_id_to_index={"req-ar": 0},
         kv_extracted_req_ids=None,
+        routed_experts_dict=None,
     )
 
     OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)

tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py

@@ -25,7 +25,11 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
 )
 
 # All tests in this file require Mooncake TransferEngine and an RDMA environment.
-pytestmark = [pytest.mark.parallel, pytest.mark.core_model]
+pytestmark = [
+    pytest.mark.parallel,
+    pytest.mark.core_model,
+    pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not installed"),
+]
 
 # ---------------------------------------------------------------------------
 # Shared helpers

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -106,18 +106,21 @@ def _build_request(
 
 def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
     """Run a single request and return (waveform, sample_rate)."""
-    # Omni.generate returns list[OmniRequestOutput] by default (py_generator=False).
-    for omni_out in omni.generate(request, DEFAULT_SAMPLING):
-        mm = omni_out.multimodal_output
-        assert mm is not None, "Expected multimodal_output to be non-None"
-        audio = mm.get("audio")
-        sr = mm.get("sr")
-        assert audio is not None, "Expected 'audio' key in multimodal_output"
-        assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
-        return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
+    for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
+        req_output = stage_outputs.request_output
+        if req_output is not None:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None, "Expected multimodal_output to be non-None"
+            audio = mm.get("audio")
+            sr = mm.get("sr")
+            assert audio is not None, "Expected 'audio' key in multimodal_output"
+            assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
+            return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
     raise AssertionError("No stage outputs received")
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     """English TTS produces non-empty 48 kHz stereo audio."""
@@ -129,6 +132,8 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)"
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Chinese TTS produces non-empty audio."""
@@ -140,6 +145,8 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
     assert not torch.all(audio == 0)
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Same seed produces identical waveforms."""
@@ -151,6 +158,8 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) ->
     assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed"
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Batch of two requests returns audio for each."""
@@ -159,11 +168,13 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
         _build_request("Second request.", ref_audio_path),
     ]
     results = []
-    # Single-stage pipeline: one SamplingParams object is broadcast to all prompts.
-    for omni_out in omni_runner.omni.generate(requests, DEFAULT_SAMPLING):
-        mm = omni_out.multimodal_output
-        assert mm is not None
-        results.append(mm["audio"].cpu())
+    # Single-stage model (num_stages=1): one sampling param for all requests.
+    for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]):
+        req_output = stage_outputs.request_output
+        if req_output is not None:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None
+            results.append(mm["audio"].cpu())
 
     assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
     for i, audio in enumerate(results):

tests/e2e/offline_inference/test_stable_audio_expansion.py

@@ -15,6 +15,7 @@ import numpy as np
 import pytest
 import torch
 
+from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.assertions import assert_audio_valid
 from tests.helpers.mark import hardware_test
 from vllm_omni import Omni
@@ -24,6 +25,8 @@ from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
+_MODEL_REPO = "stabilityai/stable-audio-open-1.0"
+
 _SAMPLE_RATE = 44100
 _CLIP_DURATION_S = 2.0
 
@@ -77,6 +80,7 @@ def test_stable_audio_quantization_and_teacache() -> None:
 
     CI should provide ``HF_TOKEN`` if the checkpoint is gated.
     """
+    skip_if_gated_repo_inaccessible(_MODEL_REPO)
     # ``model_class_name`` must be passed explicitly: the default-stage-cfg
     # factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when
     # deciding ``final_output_type`` (#2077), and at construction time the

tests/engine/test_orchestrator.py

@@ -79,7 +79,7 @@ class FakeStageClient:
         self._diffusion_outputs = queue.Queue()
 
     # Orchestrator-facing interface.
-    async def add_request_async(self, *args, **_kwargs) -> None:
+    async def add_request_async(self, *args, **kwargs) -> None:
         self.add_request_calls.append(args)
 
     async def get_output_async(self):
@@ -448,7 +448,7 @@ def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fail
             assert recursive
             return [fake_child]
 
-    def fail_base_shutdown(self):
+    def fail_base_shutdown(self, **kwargs):
         raise RuntimeError("base shutdown failed")
 
     monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
@@ -480,7 +480,7 @@ def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
 
     monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
     monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self: None)
+    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None)
 
     client = object.__new__(StageEngineCoreClient)
     client._proc = fake_proc
@@ -994,7 +994,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
 @pytest.mark.asyncio
 async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
     class FailingStageClient(FakeStageClient):
-        async def add_request_async(self, *args, **_kwargs) -> None:
+        async def add_request_async(self, *args, **kwargs) -> None:
             raise RuntimeError("submit failed")
 
     class TrackingOutputProcessor(FakeOutputProcessor):

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -389,7 +389,7 @@ class TestErrorHandling:
 
         async def empty_gen(*args, **kwargs):
             return
-            yield  # noqa: unreachable – makes this an async generator
+            yield  # unreachable – makes this an async generator
 
         engine.generate = MagicMock(side_effect=empty_gen)
         server = _make_server(engine_client=engine)
@@ -448,7 +448,7 @@ class TestErrorHandling:
 
         async def gen_value_error(*args, **kwargs):
             raise ValueError("bad value")
-            yield  # noqa: unreachable
+            yield  # unreachable
 
         engine.generate = MagicMock(side_effect=gen_value_error)
         server = _make_server(engine_client=engine)
@@ -463,7 +463,7 @@ class TestErrorHandling:
 
         async def gen_runtime_error(*args, **kwargs):
             raise RuntimeError("something went wrong")
-            yield  # noqa: unreachable
+            yield  # unreachable
 
         engine.generate = MagicMock(side_effect=gen_runtime_error)
         server = _make_server(engine_client=engine)

tests/entrypoints/openai_api/test_serving_chat_speaker.py

@@ -22,6 +22,7 @@ def serving_chat():
 
     instance = object.__new__(OmniOpenAIServingChat)
     instance._supported_speakers = None
+    instance.has_kv_connector = False
     return instance
 
 

tests/helpers/__init__.py

@@ -6,3 +6,32 @@ refactor only inside the old monolithic ``conftest``; a greedy ``__init__``
 changes import order and can affect in-process Omni (``OmniRunner`` / offline
 e2e) vs subprocess-based ``OmniServer`` tests.
 """
+
+import pytest
+
+
+def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
+    """Skip the test if a gated HuggingFace repo is not accessible.
+
+    Tries to download the model's config.json via ``hf_hub_download``,
+    which performs an actual file-access check (unlike ``HfApi().model_info()``
+    that only checks metadata).  If the token has metadata access but not
+    file-download access, ``hf_hub_download`` will raise ``GatedRepoError``
+    and we skip cleanly.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
+    except Exception:
+        return
+    try:
+        hf_hub_download(repo_id=repo_id, filename="config.json")
+    except GatedRepoError as exc:
+        pytest.skip(
+            f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
+            f"HF_TOKEN ({exc}). See docs/contributing/ci/hf_credentials.md."
+        )
+    except RepositoryNotFoundError as exc:
+        pytest.skip(f"Skipping: HF repo {repo_id!r} not found ({exc}).")
+    except Exception:
+        return

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -48,6 +48,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
     """Build the dict of modules to inject into sys.modules."""
     platforms_mock = mocker.MagicMock()
     platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
+    platforms_mock.current_omni_platform.is_npu.return_value = False
 
     logger_mock = mocker.MagicMock()
     logger_mock.init_logger = lambda name: mocker.MagicMock()
@@ -211,6 +212,10 @@ class TestCodePredictorDtypeAlignment:
         predictor._model_dtype = torch.float16
         predictor._compiled_model_fwd = predictor.model.forward
 
+        # Ensure NPU path is not taken on non-NPU hardware
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+
         # _warmup_buckets should fix the dtype mismatch
         predictor._warmup_buckets()
 
@@ -230,6 +235,9 @@ class TestCodePredictorDtypeAlignment:
         predictor = predictor.to(torch.float16)
 
         assert predictor._model_dtype is None
+        # Ensure NPU path is not taken on non-NPU hardware
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
         predictor._setup_compile()
         assert predictor._model_dtype == torch.float16
 

tests/worker/test_gpu_generation_model_runner.py

@@ -35,6 +35,7 @@ def _make_runner(multimodal_outputs):
     runner.device = torch.device("cpu")
     runner.supports_mm_inputs = False
     runner.speculative_config = None
+    runner.routed_experts_initialized = False
     return runner