[Test] Consolidate MOSS-TTS batching regression coverage

Combine overlapping row/snapshot checks, cover masked and valid attention rows in one batch, reuse the real codec session in grouping tests, and move output-policy coverage into the existing entrypoint module. Preserve both CUDA synchronization regressions and reduce the added suite from 29 to 15 cases.

Signed-off-by: Sy03 <1370724210@qq.com>

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,
     *,

tests/model_executor/models/moss_tts/test_local_output_policy.py

@@ -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"

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():

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.

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 == {}