Sy03
30d28a9a04
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>
77 lines
3.5 KiB
Python
77 lines
3.5 KiB
Python
# 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.cuda]
|
|
|
|
|
|
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() for v in values]
|
|
snapshot = _snapshot_tensor_payload_to_cpu_async(
|
|
{"list": values, "tuple": tuple(values)}, copy_stream=torch.cuda.Stream(), pin_memory=True
|
|
)
|
|
for v in values:
|
|
v.fill_(-99)
|
|
snapshot.wait()
|
|
torch.testing.assert_close(snapshot.payload, {"list": expected, "tuple": tuple(expected)}, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("offsets", [[0, 1], [1, 4]])
|
|
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
|
|
|
|
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, offsets)
|
|
finally:
|
|
torch.cuda.set_sync_debug_mode(previous)
|
|
expected = torch.full((5, 4), -10.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
|
|
|
|
|
|
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
|
|
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:
|
|
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]
|
|
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 == {}
|