215 lines
9.5 KiB
Python
215 lines
9.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.worker.test_gpu_generation_model_runner import _make_runner
|
|
from vllm_omni.worker.gpu_generation_model_runner import _AsyncGenerationOutput
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["tensor", "list", "mapping"])
|
|
def test_async_generation_owns_output_and_request_mapping(kind):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
source = torch.arange(64, device="cuda", dtype=torch.float32).reshape(2, 32)
|
|
expected = source.cpu()
|
|
raw = source if kind == "tensor" else list(source.unbind())
|
|
if kind == "mapping":
|
|
raw = {"audio": list(source.unbind())}
|
|
runner = _make_runner(raw, num_reqs=2)
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
# Publication must not wait for GPU work on the calling thread.
|
|
with (
|
|
patch.object(torch.cuda.Event, "synchronize", side_effect=AssertionError("host wait")),
|
|
patch.object(torch.cuda.Stream, "synchronize", side_effect=AssertionError("host wait")),
|
|
):
|
|
output = runner.sample_tokens()
|
|
assert isinstance(output, _AsyncGenerationOutput)
|
|
source.fill_(-1)
|
|
runner.input_batch.req_ids.clear()
|
|
runner.input_batch.req_id_to_index.clear()
|
|
resolved = output.get_output()
|
|
assert resolved.req_ids == ["req-1", "req-2"]
|
|
assert resolved.req_id_to_index == {"req-1": 0, "req-2": 1}
|
|
key = "audio" if kind == "mapping" else "model_outputs"
|
|
for i, payload in enumerate(resolved.multimodal_outputs):
|
|
torch.testing.assert_close(payload[key], expected[i])
|
|
assert output.get_output() is resolved
|
|
|
|
|
|
def test_async_generation_preserves_none_rows():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
runner = _make_runner([None, torch.ones(8, device="cuda")], num_reqs=2)
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
output = runner.sample_tokens().get_output()
|
|
assert output.multimodal_outputs[0]["model_outputs"] is None
|
|
torch.testing.assert_close(output.multimodal_outputs[1]["model_outputs"], torch.ones(8))
|
|
|
|
|
|
@pytest.mark.parametrize("owns_storage", [False, True])
|
|
def test_full_payload_accumulation_keeps_synchronous_host_copy(owns_storage):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
runner = _make_runner(torch.ones((1, 8), device="cuda"))
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=owns_storage)
|
|
runner.requests = {"req-1": object()}
|
|
seen = []
|
|
runner._should_accumulate_full_payload_output = lambda: True
|
|
runner.accumulate_full_payload_output = lambda rid, payload, state: seen.append(payload["model_outputs"].clone())
|
|
with (
|
|
patch("vllm_omni.worker.gpu_generation_model_runner.AsyncGPUModelRunnerOutput") as wrapper,
|
|
patch(
|
|
"vllm_omni.worker.gpu_generation_model_runner._snapshot_tensor_payload_to_cpu_async",
|
|
side_effect=AssertionError("unexpected deferred copy"),
|
|
),
|
|
):
|
|
runner.sample_tokens()
|
|
wrapper.assert_called_once()
|
|
torch.testing.assert_close(seen[0], torch.ones(8))
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["tensor", "list", "mapping"])
|
|
def test_owned_generation_output_skips_device_snapshot_and_retains_sources(kind):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
source = torch.empty((2, 32), device="cuda")
|
|
# Production must be ordered before copying, even if it is still pending.
|
|
torch.cuda._sleep(5_000_000)
|
|
source.fill_(7)
|
|
raw = source if kind == "tensor" else list(source.unbind())
|
|
if kind == "mapping":
|
|
raw = {"audio": list(source.unbind()), "sr": [torch.tensor(24000), torch.tensor(24000)]}
|
|
runner = _make_runner(raw, num_reqs=2)
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=True)
|
|
runner.model_config = SimpleNamespace(enforce_eager=True)
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
with (
|
|
patch(
|
|
"vllm_omni.worker.gpu_generation_model_runner._snapshot_tensor_payload_to_cpu_async",
|
|
side_effect=AssertionError("unexpected device snapshot"),
|
|
),
|
|
patch.object(torch.cuda.Event, "synchronize", side_effect=AssertionError("host wait")),
|
|
patch.object(torch.cuda.Stream, "synchronize", side_effect=AssertionError("host wait")),
|
|
):
|
|
output = runner.sample_tokens()
|
|
assert isinstance(output, _AsyncGenerationOutput)
|
|
expected_ptrs = [source.data_ptr()] if kind == "tensor" else [source[i].data_ptr() for i in range(2)]
|
|
assert [t.data_ptr() for t in output._snapshot._cuda_sources] == expected_ptrs
|
|
if kind == "mapping":
|
|
for rate in raw["sr"]:
|
|
rate.fill_(-1)
|
|
del raw, source, runner
|
|
resolved = output.get_output()
|
|
assert not output._snapshot._cuda_sources
|
|
key = "audio" if kind == "mapping" else "model_outputs"
|
|
for payload in resolved.multimodal_outputs:
|
|
torch.testing.assert_close(payload[key], torch.full((32,), 7.0))
|
|
if kind == "mapping":
|
|
assert payload["sr"].item() == 24000
|
|
assert resolved.req_ids == ["req-1", "req-2"]
|
|
assert output.get_output() is resolved
|
|
|
|
|
|
def test_owned_generation_output_preserves_none_rows():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
runner = _make_runner([None, torch.ones(8, device="cuda")], num_reqs=2)
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=True)
|
|
runner.model_config = SimpleNamespace(enforce_eager=True)
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
output = runner.sample_tokens().get_output()
|
|
assert output.multimodal_outputs[0]["model_outputs"] is None
|
|
torch.testing.assert_close(output.multimodal_outputs[1]["model_outputs"], torch.ones(8))
|
|
|
|
|
|
def test_outer_graph_execution_keeps_snapshot_even_for_owned_model():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
source = torch.ones((1, 8), device="cuda")
|
|
runner = _make_runner(source)
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=True)
|
|
runner.model_config = SimpleNamespace(enforce_eager=False)
|
|
runner.device = torch.device("cuda")
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
with patch(
|
|
"vllm_omni.worker.gpu_generation_model_runner._copy_owned_generation_payload_to_cpu_async",
|
|
side_effect=AssertionError("unsafe borrowed graph output"),
|
|
):
|
|
output = runner.sample_tokens()
|
|
source.fill_(-1)
|
|
resolved = output.get_output()
|
|
torch.testing.assert_close(resolved.multimodal_outputs[0]["model_outputs"], torch.ones(8))
|
|
|
|
|
|
@pytest.mark.parametrize("owned", [False, True])
|
|
def test_discarded_output_records_copy_stream(owned):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
source = torch.ones((2, 32), device="cuda")
|
|
runner = _make_runner(source, num_reqs=2)
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=owned)
|
|
runner.model_config = SimpleNamespace(enforce_eager=True)
|
|
runner.device = source.device
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
calls = []
|
|
original = torch.Tensor.record_stream
|
|
|
|
def record(tensor, stream):
|
|
calls.append((tensor.data_ptr(), stream))
|
|
return original(tensor, stream)
|
|
|
|
with patch.object(torch.Tensor, "record_stream", record):
|
|
output = runner.sample_tokens()
|
|
assert calls == [(output._snapshot._cuda_sources[0].data_ptr(), runner.async_output_copy_stream)]
|
|
# Drop the GPU snapshot while the copy can still be in flight. CPU storage
|
|
# remains readable after the copy stream finishes, without calling get_output.
|
|
host = output._output.multimodal_outputs
|
|
del output
|
|
runner.async_output_copy_stream.synchronize()
|
|
for row in host:
|
|
torch.testing.assert_close(row["model_outputs"], torch.ones(32))
|
|
|
|
|
|
@pytest.mark.parametrize("owned", [False, True])
|
|
@pytest.mark.parametrize("strided", [False, True])
|
|
def test_dense_waveforms_share_one_host_allocation(owned, strided):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
source = torch.arange(128, device="cuda", dtype=torch.float32).reshape(4, 32)
|
|
if strided:
|
|
source = source[:, ::2]
|
|
expected = source.cpu()
|
|
runner = _make_runner(source, num_reqs=4)
|
|
runner.model = SimpleNamespace(owns_generation_output_storage=owned)
|
|
runner.model_config = SimpleNamespace(enforce_eager=True)
|
|
runner.device = source.device
|
|
runner.use_async_scheduling = True
|
|
runner.async_output_copy_stream = torch.cuda.Stream()
|
|
output = runner.sample_tokens().get_output()
|
|
rows = [payload["model_outputs"] for payload in output.multimodal_outputs]
|
|
expected_allocations = len(rows) if strided else 1
|
|
assert len({row.untyped_storage().data_ptr() for row in rows}) == expected_allocations
|
|
assert all(row.is_contiguous() for row in rows)
|
|
for i, row in enumerate(rows):
|
|
torch.testing.assert_close(row, expected[i])
|