From b256e1ec4ee1555f7434e8b1370bfadf65f98e65 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 3 Oct 2026 22:22:39 +0800 Subject: [PATCH] Adapt Moss sampling and runner fixtures to vLLM 0.31 interfaces Signed-off-by: tzhouam --- .../sched/test_generation_batch_coalescing.py | 7 ++++- .../test_omni_ar_scheduler_kv_transfer.py | 9 +++++-- .../moss_tts/test_local_determined_tokens.py | 27 +++++++++++++------ tests/worker/test_omni_v2_profile_dispatch.py | 2 ++ tests/worker_v2/test_omni_ar_model_runner.py | 2 +- .../models/moss_tts/local_model_state.py | 18 ++++++++++--- 6 files changed, 50 insertions(+), 15 deletions(-) diff --git a/tests/core/sched/test_generation_batch_coalescing.py b/tests/core/sched/test_generation_batch_coalescing.py index 489e6a819..3cd5431b8 100644 --- a/tests/core/sched/test_generation_batch_coalescing.py +++ b/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"))]) diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py index 86102c451..a2ffb1b0a 100644 --- a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py +++ b/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, diff --git a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py index 2d529cf26..f154e7ee4 100644 --- a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py +++ b/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 diff --git a/tests/worker/test_omni_v2_profile_dispatch.py b/tests/worker/test_omni_v2_profile_dispatch.py index 949b2cace..14d86924a 100644 --- a/tests/worker/test_omni_v2_profile_dispatch.py +++ b/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 ) diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index 6bf08f147..fe89df3a2 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py index c382d4439..5a1c8f153 100644 --- a/vllm_omni/model_executor/models/moss_tts/local_model_state.py +++ b/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