[Model] Optimize MOSS-TTS Local batched execution and streaming codec

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

tests/worker/test_batched_omni_output.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+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]
+
+
+@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):
+    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
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        {"codes": payload, "metadata": "kept"}, 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)
+
+
+@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