[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

@@ -7,7 +7,11 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
+from vllm.config import VllmConfig
 from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
+from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
+from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.states import RequestState
 
@@ -27,11 +31,11 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed,
         pytest.skip("CUDA required")
     device = torch.device("cuda")
     reqs = RequestState(5, 32, 16, 0, 8, device)
-    sampler = Sampler(5, 8, device, reqs)
+    sampler = Sampler(VllmConfig(), 5, 8, device, reqs)
     for i, length in enumerate([4, 1, 3]):
         reqs.add_request(str(i), length, [2] * length, 0, 8)
         slot = reqs.req_id_to_index[str(i)]
-        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+        sampler.add_request(slot, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
     reqs.apply_staged_writes()
     sampler.apply_staged_writes()
     state = _state(MossLocalModelState, device)
@@ -39,6 +43,7 @@ def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed,
     state.model._batch_state = None
     batch = _batch(device, [2, 4, 3], [1, 2, 1])
     batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
+    batch.seq_lens_cpu_upper_bound = torch.tensor([5, 2, 2], dtype=torch.int32)
     batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
     batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
     batch.expanded_idx_mapping = batch.idx_mapping
@@ -73,10 +78,10 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint):
         pytest.skip("CUDA required")
     device = torch.device("cuda")
     reqs = RequestState(5, 32, 16, 0, 8, device)
-    sampler = Sampler(5, 8, device, reqs)
+    sampler = Sampler(VllmConfig(), 5, 8, device, reqs)
     reqs.add_request("constrained", 1, [2], 0, 8)
     slot = reqs.req_id_to_index["constrained"]
-    sampler.add_request(slot, 1, SamplingParams(**constraint))
+    sampler.add_request(slot, SamplingParams(**constraint))
     state = _state(MossLocalModelState, device)
     state._direct_tokens = True
     # Use the actual upstream request registration, including its combined
@@ -85,16 +90,22 @@ def test_real_sampler_constraints_disable_determined_tokens(constraint):
 
 
 @pytest.mark.cpu
-@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
+@pytest.mark.parametrize(
+    "feature", ["logprobs", "bias", "penalty", "bad_words", "custom", "thinking", "trace", "mask", "nans"]
+)
 def test_distribution_features_fall_back_before_using_gpu(feature):
     sampler = object.__new__(Sampler)
     sampler.compute_nans = feature == "nans"
     sampler.return_sampling_mask = feature == "mask"
     sampler.trace_replay_state = object() if feature == "trace" else None
     sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
-    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
-    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
-    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
+    bias = object.__new__(LogitBiasState)
+    bias.use_logit_bias = np.array([feature == "bias"])
+    penalties = object.__new__(PenaltiesState)
+    penalties.use_penalty = np.array([feature == "penalty"])
+    bad_words = object.__new__(BadWordsState)
+    bad_words.num_bad_words = SimpleNamespace(np=np.array([feature == "bad_words"]))
+    sampler.logits_processors = [bias, penalties, bad_words] + ([object()] if feature == "custom" else [])
     sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
     state = object.__new__(MossLocalModelState)
     state._direct_tokens = True

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

@@ -132,10 +132,10 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc
         hidden_states=torch.zeros(2, 4),
         finished_req_ids=set(),
         ec_connector_output=None,
-        routed_experts=None,
     )
     runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
     runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.aux_output_connector = None
     runner.model_config = SimpleNamespace(async_chunk=False)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
     runner.model_state, runner.model = state, model
@@ -154,7 +154,10 @@ def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatc
             torch.zeros(1),
         )
     )
-    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
+    runner.prompt_logprobs_worker = SimpleNamespace(
+        compute_prompt_logprobs=mocker.Mock(return_value={}),
+        compute_prompt_token_id_logprobs=mocker.Mock(return_value={}),
+    )
     runner.postprocess_sampled = mocker.Mock()
     runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
     mock_out = SimpleNamespace(copy_event=None)

tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py

@@ -42,10 +42,26 @@ def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProce
     return processor
 
 
+class _FakeProcessorItem:
+    def __init__(self, processor_data: dict[str, object]) -> None:
+        self._processor_data = processor_data
+
+    def get_processor_data(self) -> dict[str, object]:
+        return dict(self._processor_data)
+
+    def get_passthrough_data(self) -> dict[str, object]:
+        return {}
+
+
 def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace:
+    class _Selected:
+        @staticmethod
+        def values():
+            return [_FakeProcessorItem(mm_data)] if mm_data else []
+
     return SimpleNamespace(
         get_all_counts=lambda: counts,
-        select=lambda _keys: mm_data,
+        select=lambda _keys: _Selected(),
     )
 
 
@@ -75,9 +91,8 @@ def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> Non
 def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None:
     processor = _make_processor(monkeypatch)
     audio = np.zeros(2400, dtype=np.float32)
-    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {}))
 
-    out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"})
+    out = processor._apply_hf_processor_main(_items({"audio": 1}, {"audios": [audio]}), {"ref_text": "reference"})
 
     assert tuple(out["ref_audio_tokens"].shape) == (8, 4)
     assert list(out["ref_audio_len"]) == [4]
@@ -89,7 +104,6 @@ def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.Monke
 def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None:
     """Text-only requests and fully cached audio both reach this path with empty items."""
     processor = _make_processor(monkeypatch)
-    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {}))
 
     out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"})
 
@@ -99,4 +113,6 @@ def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.Monke
 
 
 def test_legacy_call_hf_processor_is_gone() -> None:
-    assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor")
+    # Upstream reintroduced a 2-arg `_call_hf_processor` on the base class; the
+    # omni-specific 4-arg hook must not be defined on the omni class itself.
+    assert "_call_hf_processor" not in OmniVoiceMultiModalProcessor.__dict__

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -262,6 +262,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         "audio": 1,
     }[modality]
     mm_data_items.__getitem__ = MagicMock(side_effect=lambda modality: [f"{modality}-0"])
+    mm_data_items.values = MagicMock(return_value=[])
 
     inputs = MagicMock()
     inputs.mm_data_items = mm_data_items
@@ -275,6 +276,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     cache = MagicMock()
     # Video miss, audio hit under pair-aware keys — expansion must force audio.
     cache.is_cached.side_effect = lambda keys: [False] if keys[0].startswith("v") else [True]
+    inputs.cache = cache
 
     captured = {}
 
@@ -309,6 +311,8 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     fake_self._get_mm_fields_config = MagicMock(return_value={})
     fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]})
     fake_self._merge_mm_kwargs = capture_merge
+    fake_self._postprocess_prompt = lambda prompt: prompt
+    inputs.prompt = [0, 1, 2]
 
     with patch(
         "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
@@ -549,24 +553,6 @@ def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count():
         )
 
 
-def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio():
-    captured_counts = {}
-    fake_self = SimpleNamespace(
-        dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"),
-        _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False),
-    )
-    mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1})
-
-    Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only(
-        fake_self,
-        mm_items,
-        {"use_audio_in_video": [True, False]},
-        {},
-    )
-
-    assert captured_counts == {"video": 2, "audio": 0}
-
-
 def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits():
     fake_self = _fake_processor()
     update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID]))
@@ -661,7 +647,12 @@ def test_qwen3_processor_inherits_vllm_omni_per_video_helpers():
         Qwen3OmniMoeThinkerMultiModalProcessor,
         Qwen2_5OmniThinkerMultiModalProcessor,
     )
-    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data
+    # Keep Qwen3's native HF preprocessing while inheriting Omni pair handling.
+    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_inputs is UpstreamQwen3Processor._get_hf_mm_inputs
+    assert (
+        Qwen3OmniMoeThinkerMultiModalProcessor._get_audio_in_video_pairs
+        is Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs
+    )
     mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__
     assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor)
 

tests/model_executor/models/test_omni_processing.py

@@ -79,31 +79,16 @@ def get_text_token_prompts(
     """Return ``(text_prompt, token_prompt)`` for the given multimodal data."""
     dummy_inputs = processor.dummy_inputs
     tokenizer: TokenizerLike = processor.info.get_tokenizer()
-    model_config = processor.info.ctx.model_config
 
     parsed_data = processor.info.parse_mm_data(mm_data)
     mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
 
-    inputs = dummy_inputs.get_dummy_processor_inputs(
-        model_config.max_model_len,
-        mm_counts,
-        mm_options={},
+    text_prompt = dummy_inputs.get_dummy_text(mm_counts)
+    token_prompt = tokenizer.encode(
+        text_prompt,
+        **processor.info.get_default_tok_params().get_encode_kwargs(),
     )
 
-    text_prompt: str | None
-    token_prompt: list[int]
-    if isinstance(inputs.prompt, list):
-        text_prompt = None
-        token_prompt = inputs.prompt
-    elif isinstance(inputs.prompt, str):
-        text_prompt = inputs.prompt
-        token_prompt = tokenizer.encode(
-            text_prompt,
-            **processor.info.get_default_tok_params().get_encode_kwargs(),
-        )
-    else:
-        raise TypeError(type(inputs.prompt))
-
     return text_prompt, token_prompt
 
 
@@ -199,8 +184,8 @@ def _test_processing_correctness(
         modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items()
     }
 
-    baseline_processor = factories.build_processor(ctx, cache=None)
-    cached_processor = factories.build_processor(ctx, cache=cache)
+    baseline_processor = factories.build_processor(ctx)
+    cached_processor = factories.build_processor(ctx)
 
     rng = np.random.RandomState(0)
 
@@ -238,6 +223,7 @@ def _test_processing_correctness(
             baseline_processor,
             cached_processor,
             batch_idx,
+            cache,
         )
 
 
@@ -247,6 +233,7 @@ def _test_processing_correctness_one(
     baseline_processor: BaseMultiModalProcessor,
     cached_processor: BaseMultiModalProcessor,
     batch_idx: int,
+    cache: MultiModalProcessorOnlyCache,
 ):
     model_type = model_config.hf_config.model_type
 
@@ -258,12 +245,14 @@ def _test_processing_correctness_one(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=None,
     )
 
     cached_tokenized_result = cached_processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
 
     _assert_inputs_equal(
@@ -278,11 +267,13 @@ def _test_processing_correctness_one(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
+            cache=None,
         )
         cached_text_result = cached_processor(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
+            cache=cache,
         )
 
         _assert_inputs_equal(