vllm-omni/tests/worker/test_batched_omni_output.py

# 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


@pytest.mark.cuda
def test_mtp_mixed_offsets_do_not_synchronize_cuda(monkeypatch):
    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")}
    for name in ("talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"):
        buffer = getattr(runner, name)
        buffer.gpu = buffer.gpu.to("cuda")
    embeds = torch.ones((2, 4), device="cuda")
    codes = torch.arange(2, device="cuda").view(2, 1)
    runner.talker_mtp = lambda *args, **kwargs: (embeds, codes)
    output = torch.full((5, 4), -10.0, device="cuda")
    previous = torch.cuda.get_sync_debug_mode()
    try:
        torch.cuda.set_sync_debug_mode("error")
        runner._talker_mtp_forward(["r1", "r2"], output, [1, 4])
    finally:
        torch.cuda.set_sync_debug_mode(previous)
    expected = torch.full((5, 4), -10.0)
    expected[[1, 4]] = 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)}},
    ]
    hidden = torch.ones((3, 4), device="cuda")
    previous = torch.cuda.get_sync_debug_mode()
    try:
        torch.cuda.set_sync_debug_mode("error")
        output = model.make_omni_output(hidden, runtime_additional_information=infos)
    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)]