Adapt Moss sampling and runner fixtures to vLLM 0.31 interfaces

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

tests/core/sched/test_generation_batch_coalescing.py

@@ -171,6 +171,7 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s, _ = scheduler(monkeypatch, [event])
     s._generation_max_wait_s = 0
     s.waiting = []
+    s.kv_holding_waiting = []
     s.input_coordinator._async_chunk = True
     s.input_coordinator.update_request_metadata = mocker.Mock()
     s.input_coordinator.process_pending_chunks = mocker.Mock()
@@ -178,7 +179,11 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s.input_coordinator.update_request_metadata.assert_called_once_with(
         s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
     )
-    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+    s.input_coordinator.process_pending_chunks.assert_called_once()
+    waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args
+    assert list(waiting) == []
+    assert running is s.running
+    assert ready == {"0"} and finished == set()
 
 
 @pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -185,10 +185,15 @@ def test_stage_zero_final_skips_inter_stage_output(mocker):
 def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler._omits_kv_transfer_cache = {}
-    scheduler.skipped_waiting = []
+    scheduler.deferred_waiting = set()
+    scheduler.kv_holding_waiting = []
     scheduler.log_stats = False
     tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
-    session = SimpleNamespace(
+
+    class Session(SimpleNamespace):
+        __hash__ = object.__hash__
+
+    session = Session(
         request_id="req",
         additional_information=tagged.additional_information,
         status=RequestStatus.WAITING,

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/worker/test_omni_v2_profile_dispatch.py

@@ -16,6 +16,7 @@ def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
     runner = object.__new__(OmniGPUModelRunner)
     runner.cudagraph_manager = None
     runner.parallel_config = SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=1)
+    runner.dp_size, runner.dp_rank = dp_size, 1
 
     import vllm.v1.worker.gpu.dp_utils as dp_utils
 
@@ -52,6 +53,7 @@ def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
 
     runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
     runner.parallel_config.data_parallel_size = 1
+    runner.dp_size = 1
     batch_desc, _ = runner._dispatch_batch_descriptor(
         num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8
     )

tests/worker_v2/test_omni_ar_model_runner.py

@@ -804,7 +804,7 @@ def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocke
     extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None
     result = _async_output(
         req_ids=["frame", "eos", "prefill"],
-        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None),
+        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None, None),
         num_sampled_tokens=torch.tensor([1, 1, 1]),
         multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}},
         input_batch=batch,

vllm_omni/model_executor/models/moss_tts/local_model_state.py

@@ -14,7 +14,10 @@ from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected
+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.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
@@ -118,15 +121,24 @@ class MossLocalModelState(OmniModelState):
             or sampler.return_sampling_mask
             or sampler.trace_replay_state is not None
             or sampler.get_logprobs_dims(rows) is not None
-            or np.any(sampler.logit_bias_state.use_logit_bias[rows])
-            or np.any(sampler.penalties_state.use_penalty[rows])
-            or np.any(sampler.bad_words_state.num_bad_words.np[rows])
             or (
                 sampler.thinking_budget_state.enabled
                 and np.any(sampler.thinking_budget_state.use_thinking_budget[rows])
             )
         ):
             return None
+        for processor in sampler.logits_processors:
+            if type(processor) is LogitBiasState:
+                active = np.any(processor.use_logit_bias[rows])
+            elif type(processor) is PenaltiesState:
+                active = np.any(processor.use_penalty[rows])
+            elif type(processor) is BadWordsState:
+                active = np.any(processor.num_bad_words.np[rows])
+            else:
+                # Custom processors can alter even a single finite logit.
+                return None
+            if active:
+                return None
         keep = self.model._batch_should_continue
         if keep is None or keep.numel() != batch.num_reqs:
             return None