diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index a93ae0dbf..6fe04fdb7 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -227,6 +227,25 @@ def test_resolve_sampling_params_list_preserves_stage_constraints(): assert caller_params.stop_token_ids == [7] +@pytest.mark.parametrize("use_defaults", [False, True]) +def test_moss_local_output_policy_preserves_codec_streaming(use_defaults): + from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE + + base = _make_base() + base.engine.num_stages = 2 + base.sampling_constraints_list = [stage.sampling_constraints for stage in MOSS_TTS_LOCAL_PIPELINE.stages] + base.default_sampling_params_list = [ + base._apply_sampling_constraints(SamplingParams(), constraints) + for constraints in base.sampling_constraints_list + ] + caller = [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(2)] + result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True) + + assert [params.output_kind for params in result] == [RequestOutputKind.FINAL_ONLY, RequestOutputKind.DELTA] + assert all(params.output_kind == RequestOutputKind.DELTA for params in caller) + assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY + + def _stage_spec( stage_id: int, *, diff --git a/tests/model_executor/models/moss_tts/test_local_output_policy.py b/tests/model_executor/models/moss_tts/test_local_output_policy.py deleted file mode 100644 index 4c367c049..000000000 --- a/tests/model_executor/models/moss_tts/test_local_output_policy.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest -from vllm import SamplingParams -from vllm.sampling_params import RequestOutputKind - -from vllm_omni.entrypoints.omni_base import OmniBase -from vllm_omni.entrypoints.utils import coerce_param_message_types -from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.mark.parametrize("use_defaults", [False, True]) -def test_local_talker_control_outputs_preserve_codec_streaming(mocker, use_defaults): - base = object.__new__(OmniBase) - base.engine = mocker.Mock(num_stages=2) - base.sampling_constraints_list = [dict(stage.sampling_constraints) for stage in MOSS_TTS_LOCAL_PIPELINE.stages] - base.default_sampling_params_list = [ - base._apply_sampling_constraints(SamplingParams(max_tokens=8, seed=11), constraints) - for constraints in base.sampling_constraints_list - ] - caller = coerce_param_message_types( - [SamplingParams(max_tokens=8, seed=11), SamplingParams(max_tokens=8, seed=11)], True - ) - result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True) - - assert result[0].output_kind == RequestOutputKind.FINAL_ONLY - assert result[0].detokenize is False - assert result[0].max_tokens == 8 - assert result[0].seed == 11 - assert result[1].output_kind == RequestOutputKind.DELTA - assert caller[0].output_kind == RequestOutputKind.DELTA - assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY - assert MOSS_TTS_LOCAL_PIPELINE.stages[0].final_output is False - assert MOSS_TTS_LOCAL_PIPELINE.stages[1].final_output_type == "audio" diff --git a/tests/model_executor/models/moss_tts/test_streaming_attention.py b/tests/model_executor/models/moss_tts/test_streaming_attention.py index fc77232a6..3134575ca 100644 --- a/tests/model_executor/models/moss_tts/test_streaming_attention.py +++ b/tests/model_executor/models/moss_tts/test_streaming_attention.py @@ -16,10 +16,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cuda] @pytest.mark.parametrize( - "batch,heads,queries,keys", [(1, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)] + "batch,heads,queries,keys", [(2, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)] ) -@pytest.mark.parametrize("empty_mask", [False, True]) -def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys, empty_mask): +def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys): torch.manual_seed(7) q = torch.randn(batch, queries, 3, heads, 64, device="cuda", dtype=torch.bfloat16)[:, :, 0].permute(0, 2, 1, 3) k = torch.randn(batch, heads, keys, 64, device="cuda", dtype=torch.bfloat16) @@ -27,14 +26,12 @@ def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, head positions = torch.arange(keys, device="cuda").roll(17) query_positions = torch.arange(queries, device="cuda") + keys - queries mask = (positions[None, :] <= query_positions[:, None])[None, None].expand(batch, 1, -1, -1).clone() - if empty_mask: - mask.zero_() + mask[0, :, 0] = False # Cover an all-masked row alongside valid rows. actual = masked_attention(q, k, v, mask) expected = F.scaled_dot_product_attention(q, k, v, mask) assert bool(actual.isfinite().all()) torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01) - if empty_mask: - assert actual.count_nonzero() == 0 + assert actual[0, :, 0].count_nonzero() == 0 def test_streaming_attention_preserves_ring_wrap_reordering_and_slot_reset(): diff --git a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py index 585c3da49..fff97cd96 100644 --- a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py +++ b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py @@ -6,7 +6,10 @@ import pytest import torch from torch import nn -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import _MossCodecStreamSession +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import ( + MossTTSCodecDecoder, + _MossCodecStreamSession, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -31,77 +34,45 @@ class CausalCodec(nn.Module): def session(): - return _MossCodecStreamSession(CausalCodec(), state_capacity=3, n_vq=2, vllm_config=None) + return _MossCodecStreamSession(CausalCodec(), state_capacity=8, n_vq=2, vllm_config=None) -@pytest.mark.parametrize("tail_length", [1, 3, 7, 14]) -def test_terminal_padding_preserves_audio_live_state_and_reused_slots(tail_length): - batched, separate = session(), session() - for s in (batched, separate): - assert [s.acquire(), s.acquire()] == [0, 1] - s.step({0: torch.ones(2, 15), 1: torch.full((2, 15), 3)}) - short = torch.arange(2 * tail_length).reshape(2, tail_length) - long = torch.full((2, 15), 5) - expected_short = separate.step({0: short}, terminal_slots={0})[0] - expected_long = separate.step({1: long})[1] - actual = batched.step({0: short, 1: long}, terminal_slots={0}) - torch.testing.assert_close(actual[0], expected_short, rtol=0, atol=0) - torch.testing.assert_close(actual[1], expected_long, rtol=0, atol=0) - for s in (batched, separate): - s.release(0, state_already_reset=True) - assert s.acquire() == 0 - next_step = {0: torch.ones(2, 1), 1: torch.ones(2, 1)} - a, b = batched.step(next_step), separate.step(next_step) - for slot in (0, 1): - torch.testing.assert_close(a[slot], b[slot], rtol=0, atol=0) - - -def test_nonterminal_row_cannot_silently_advance_through_padding(): +def test_terminal_padding_preserves_audio_live_state_and_reused_slots(): s = session() - s.acquire() - s.acquire() + assert [s.acquire(), s.acquire()] == [0, 1] + plan = {0: torch.ones(2, 3), 1: torch.ones(2, 15)} with pytest.raises(ValueError, match="Only terminal"): - s.step({0: torch.ones(2, 3), 1: torch.ones(2, 15)}) + s.step(plan) assert s._codec.state.count_nonzero() == 0 + s.step({0: torch.ones(2, 1), 1: torch.full((2, 1), 2)}) + actual = s.step(plan, terminal_slots={0}) + torch.testing.assert_close(actual[0], torch.tensor([[4.0, 6.0, 8.0]])) + torch.testing.assert_close(actual[1], torch.arange(6.0, 36.0, 2.0)[None]) + s.release(0, state_already_reset=True) + assert s.acquire() == 0 + resumed = s.step({0: torch.ones(2, 1), 1: torch.ones(2, 1)}) + assert [resumed[i].item() for i in (0, 1)] == [2, 36] @pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)]) -def test_only_existing_compatible_graph_buckets_are_coalesced(graph_capacity, expected_calls): - from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder - - calls = [] - leased: list[int] = [] - - def acquire(): - slot = len(leased) - leased.append(slot) - return slot - - def step(plan, *, terminal_slots): - calls.append((dict(plan), set(terminal_slots))) - return {slot: torch.zeros(1, codes.shape[1]) for slot, codes in plan.items()} - - graph = ( +def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls): + s = session() + s._cudagraph_wrapper = ( None if graph_capacity is None else SimpleNamespace( batch_sizes=[1, graph_capacity], - frame_sizes=[1, 8, 15], _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None), + decode=lambda *args, **kwargs: None, ) ) - fake_session = SimpleNamespace( - _cudagraph_wrapper=graph, - acquire=acquire, - step=step, - release=lambda slot, **kwargs: None, - ) + step = mocker.spy(s, "step") decoder = object.__new__(MossTTSCodecDecoder) nn.Module.__init__(decoder) decoder._stream_max_step_frames = 15 decoder._stream_state_capacity = 8 decoder._stream_req_slots = {} - decoder._ensure_stream_session = lambda: fake_session + decoder._ensure_stream_session = lambda: s lengths = [3, 7, 8, 1, 15] finished = [True, True, False, True, False] result = decoder._decode_streaming_batch( @@ -110,9 +81,9 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(graph_capacity, ex for i, (length, done) in enumerate(zip(lengths, finished, strict=True)) ] ) - assert len(calls) == expected_calls + assert step.call_count == expected_calls assert [result[i].shape[-1] for i in range(5)] == lengths if graph_capacity == 8: - assert set(calls[0][0]) == {0, 1, 2} - assert calls[0][1] == {0, 1} - assert set(calls[1][0]) == {3} # Preserve the exact one-frame graph. + assert list(step.call_args_list[0].args[0]) == [0, 1, 2] + assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1} + assert list(step.call_args_list[1].args[0]) == [3] # Preserve the exact one-frame graph. diff --git a/tests/worker/test_batched_omni_output.py b/tests/worker/test_batched_omni_output.py index b9b86991a..d0b18b3c1 100644 --- a/tests/worker/test_batched_omni_output.py +++ b/tests/worker/test_batched_omni_output.py @@ -6,96 +6,27 @@ import torch from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] -@pytest.mark.cpu -@pytest.mark.parametrize("kinds", [("emit", "stop", "prefill"), ("prefill", "emit"), ("stop",), ("prefill",)]) -def test_local_stop_flags_and_output_rows_stay_request_aligned(kinds): - model = object.__new__(MossTTSLocalTalkerForGeneration) - torch.nn.Module.__init__(model) - model.n_vq = 12 - model.audio_pad_token_id = 1024 - infos = [] - originals = [] - for index, kind in enumerate(kinds): - codes = torch.full((1, 12), 1024 if kind == "stop" else index, dtype=torch.long) - originals.append(codes) - infos.append({} if kind == "prefill" else {"audio_codes": {"current": codes}}) - result = model.make_omni_output(torch.ones((len(kinds), 4)), runtime_additional_information=infos) - if all(kind == "prefill" for kind in kinds): - assert model._batch_should_continue is None - assert result.multimodal_outputs == {} - return - assert model._batch_should_continue.tolist() == [kind != "stop" for kind in kinds] - rows = result.multimodal_outputs["codes"]["audio"] - assert len(rows) == len(kinds) - for index, kind in enumerate(kinds): - if kind == "prefill": - assert rows[index].shape == (0, 12) - else: - torch.testing.assert_close(rows[index], originals[index], rtol=0, atol=0) - assert all("current" not in info.get("audio_codes", {}) for info in infos) - - -@pytest.mark.cuda -@pytest.mark.parametrize("as_tuple", [False, True]) -def test_async_packed_snapshot_survives_source_reuse(as_tuple): +def test_async_packed_snapshot_survives_source_reuse(): values = [ torch.arange(12, device="cuda").view(1, 12), torch.empty((0, 12), device="cuda", dtype=torch.long), torch.arange(24, device="cuda").view(2, 12), ] - expected = [v.cpu().clone() for v in values] - payload = tuple(values) if as_tuple else values + expected = [v.cpu() for v in values] snapshot = _snapshot_tensor_payload_to_cpu_async( - {"codes": payload, "metadata": "kept"}, copy_stream=torch.cuda.Stream(), pin_memory=True + {"list": values, "tuple": tuple(values)}, copy_stream=torch.cuda.Stream(), pin_memory=True ) for v in values: v.fill_(-99) snapshot.wait() - actual = snapshot.payload["codes"] - assert isinstance(actual, tuple if as_tuple else list) - assert snapshot.payload["metadata"] == "kept" - for out, ref in zip(actual, expected, strict=True): - torch.testing.assert_close(out, ref, rtol=0, atol=0) - # A second output must not alias the first output's host snapshot either. - later = _snapshot_tensor_payload_to_cpu_async(values, copy_stream=torch.cuda.Stream(), pin_memory=True) - later.wait() - for out, ref in zip(actual, expected, strict=True): - torch.testing.assert_close(out, ref, rtol=0, atol=0) + torch.testing.assert_close(snapshot.payload, {"list": expected, "tuple": tuple(expected)}, rtol=0, atol=0) -@pytest.mark.cpu @pytest.mark.parametrize("offsets", [[0, 1], [1, 4]]) -def test_owned_code_batch_and_embeddings_survive_next_graph_output(monkeypatch, offsets): - import vllm_omni.worker.gpu_model_runner as mod - from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context - - monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context) - runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) - runner.model.gpu_resident_buffer_keys = {("codes", "audio")} - original = runner.talker_mtp - returned = [] - - def mtp(*args, **kwargs): - result = original(*args, **kwargs) - returned.append(result) - return result - - runner.talker_mtp = mtp - embeddings = torch.full((6, 4), -10.0) - runner._talker_mtp_forward(["r1", "r2"], embeddings, offsets) - assert torch.equal(embeddings[offsets], torch.ones(2, 4)) - for i in set(range(6)) - set(offsets): - assert torch.equal(embeddings[i], torch.full((4,), -10.0)) - returned[0][1].fill_(999) - assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0 - assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1 - - -@pytest.mark.cuda -def test_mtp_mixed_offsets_do_not_synchronize_cuda(monkeypatch): +def test_mtp_preserves_rows_and_snapshot_without_cuda_sync(monkeypatch, offsets): import vllm_omni.worker.gpu_model_runner as mod from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context @@ -112,28 +43,25 @@ def test_mtp_mixed_offsets_do_not_synchronize_cuda(monkeypatch): previous = torch.cuda.get_sync_debug_mode() try: torch.cuda.set_sync_debug_mode("error") - runner._talker_mtp_forward(["r1", "r2"], output, [1, 4]) + runner._talker_mtp_forward(["r1", "r2"], output, offsets) finally: torch.cuda.set_sync_debug_mode(previous) expected = torch.full((5, 4), -10.0) - expected[[1, 4]] = 1.0 + expected[offsets] = 1.0 torch.testing.assert_close(output.cpu(), expected, rtol=0, atol=0) codes.fill_(999) assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0 assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1 -@pytest.mark.cuda def test_mixed_stop_flags_do_not_synchronize_cuda(): model = object.__new__(MossTTSLocalTalkerForGeneration) torch.nn.Module.__init__(model) model.n_vq = 12 model.audio_pad_token_id = 1024 - infos = [ - {"audio_codes": {"current": torch.zeros((1, 12), device="cuda", dtype=torch.long)}}, - {}, - {"audio_codes": {"current": torch.full((1, 12), 1024, device="cuda", dtype=torch.long)}}, - ] + valid = torch.arange(12, device="cuda").view(1, 12) + stopped = torch.full_like(valid, 1024) + infos = [{"audio_codes": {"current": valid}}, {}, {"audio_codes": {"current": stopped}}] hidden = torch.ones((3, 4), device="cuda") previous = torch.cuda.get_sync_debug_mode() try: @@ -142,4 +70,7 @@ def test_mixed_stop_flags_do_not_synchronize_cuda(): finally: torch.cuda.set_sync_debug_mode(previous) assert model._batch_should_continue.tolist() == [True, True, False] - assert [row.shape for row in output.multimodal_outputs["codes"]["audio"]] == [(1, 12), (0, 12), (1, 12)] + torch.testing.assert_close(output.multimodal_outputs["codes"]["audio"], [valid, valid[:0], stopped]) + assert all("current" not in info.get("audio_codes", {}) for info in infos) + prefill = model.make_omni_output(hidden[:1], runtime_additional_information=[{}]) + assert model._batch_should_continue is None and prefill.multimodal_outputs == {}