diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 9bbd6c323..bb350a3c6 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -122,7 +122,7 @@ steps: commands: - | timeout 15m bash -c " - pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model' + pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker/test_gpu_generation_async_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model' " mirror_hardwares: l4_1 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index bb6229507..b7ada6ecd 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -1109,3 +1109,14 @@ def test_request_conditioning_reaches_code2wav_with_true_prompt_length(): assert call["prompt_token"].shape == (1, 2) assert call["prompt_feat"].shape == (1, 4, 80) assert output.multimodal_outputs + + +@pytest.mark.parametrize("stage", ["cosyvoice3_talker", "cosyvoice3_code2wav"]) +@pytest.mark.parametrize("owned", [False, True]) +def test_generation_output_ownership_is_exposed_only_for_owned_codec(stage, owned): + model = object.__new__(CosyVoice3Model) + nn.Module.__init__(model) + model.model_stage = stage + if stage == "cosyvoice3_code2wav": + model.code2wav = SimpleNamespace(owns_generation_output_storage=owned) + assert model.owns_generation_output_storage is (stage == "cosyvoice3_code2wav" and owned) diff --git a/tests/worker/test_gpu_generation_async_output.py b/tests/worker/test_gpu_generation_async_output.py new file mode 100644 index 000000000..93481be3b --- /dev/null +++ b/tests/worker/test_gpu_generation_async_output.py @@ -0,0 +1,214 @@ +# 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]) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index ba96348e0..e5b62e9a5 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -744,6 +744,13 @@ class CosyVoice3Model( prefer_model_sampler = True _sampling_eps = 1e-5 + @property + def owns_generation_output_storage(self) -> bool: + """Expose the codec's ownership contract on the model seen by the runner.""" + return self.model_stage == "cosyvoice3_code2wav" and ( + getattr(self.code2wav, "owns_generation_output_storage", False) is True + ) + @property def logits_vocab_size(self) -> int | None: """Width of the logits ``compute_logits`` returns (the speech head). diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 96ec400c6..96177549c 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -92,6 +92,11 @@ class CosyVoice3Code2Wav(nn.Module): - HiFiGAN vocoder for mel-to-waveform conversion """ + # Returned CUDA payloads own their storage (including HiFT graph results). + # Subsequent forwards/state updates never mutate it; the runner can retain + # the tensors and copy to host without taking another device snapshot. + owns_generation_output_storage = True + def __init__(self, config: CosyVoice3Config): super().__init__() self.config = config diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 25a69a1de..85bd43f3d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -270,6 +270,11 @@ class MiniCPMO45Code2Wav(nn.Module): has_preprocess = False has_postprocess = False + # Returned CUDA payloads own their storage (including HiFT graph results). + # Subsequent forwards/state updates never mutate it; the runner can retain + # the tensors and copy to host without taking another device snapshot. + owns_generation_output_storage = True + def __init__( self, *, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 78e653b54..36db21065 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -182,6 +182,10 @@ def _snapshot_tensor_payload_to_cpu_async( with torch.cuda.stream(copy_stream): copy_stream.wait_stream(source_stream) cpu_payload = _copy_tensor_payload_to_cpu(cloned, pin_memory) + # An executor may discard an unconsumed output (e.g. cancellation). + # Keep the allocator from reusing its source storage before D2H ends. + for source in cuda_sources: + source.record_stream(copy_stream) ready_event.record(copy_stream) return _AsyncCPUPayloadSnapshot(cpu_payload, ready_event, cuda_sources) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 7c2cf2163..216d31f44 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -13,7 +13,7 @@ import gc import logging from collections.abc import Mapping from dataclasses import replace -from typing import cast +from typing import Any, cast import numpy as np import torch @@ -42,7 +42,13 @@ from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_payload_list -from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values +from vllm_omni.worker.gpu_ar_model_runner import ( + ExecuteModelState, + _AsyncCPUPayloadSnapshot, + _copy_tensor_payload_to_cpu, + _ensure_tensor_values, + _snapshot_tensor_payload_to_cpu_async, +) from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.mixins import maybe_unpad_input_ids from vllm_omni.worker.omni_connector_model_runner_mixin import ( @@ -53,6 +59,57 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import ( logger = logging.getLogger(__name__) +class _AsyncGenerationOutput(AsyncModelRunnerOutput): + """Generation payloads become readable when their copy-stream event completes.""" + + def __init__(self, output: OmniModelRunnerOutput, snapshot: _AsyncCPUPayloadSnapshot) -> None: + self._output = output + self._snapshot = snapshot + + def get_output(self) -> OmniModelRunnerOutput: + self._snapshot.wait() + return self._output + + +def _copy_owned_generation_payload_to_cpu_async( + payload: Any, *, copy_stream: torch.cuda.Stream, pin_memory: bool +) -> _AsyncCPUPayloadSnapshot: + """Copy immutable model-owned outputs, retaining their storage until D2H finishes. + + The model must guarantee no subsequent forward or state update mutates any + returned CUDA payload storage. CPU values still need a synchronous snapshot. + record_stream also protects allocations if an unconsumed output is dropped. + """ + sources: list[torch.Tensor] = [] + + def retain(value: Any) -> Any: + if isinstance(value, torch.Tensor): + if value.device.type == "cuda": + tensor = value.detach() + sources.append(tensor) + return tensor + return value.detach().clone() + if isinstance(value, dict): + return {key: retain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + items = [retain(item) for item in value] + return tuple(items) if isinstance(value, tuple) else items + return value + + retained = retain(payload) + if not sources: + return _AsyncCPUPayloadSnapshot(retained, None, sources) + producer = torch.cuda.current_stream() + ready = torch.cuda.Event() + with torch.cuda.stream(copy_stream): + copy_stream.wait_stream(producer) + host_payload = _copy_tensor_payload_to_cpu(retained, pin_memory) + for source in sources: + source.record_stream(copy_stream) + ready.record(copy_stream) + return _AsyncCPUPayloadSnapshot(host_payload, ready, sources) + + class _HostCopyBatch: """Device-to-host copies of one step's outputs behind a single host sync. @@ -442,8 +499,15 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # Build per-request multimodal_outputs list (dedicated channel). # pooler_output is no longer used for multimodal data. per_req_payloads: list[dict[str, object]] = [] + pin_memory = is_pin_memory_available() + accumulate_full_payload = self._should_accumulate_full_payload_output() + use_async_output = ( + self.use_async_scheduling and self.device.type == "cuda" and pin_memory and not accumulate_full_payload + ) # One host sync for the whole step instead of one per request tensor. - to_host = _HostCopyBatch(is_pin_memory_available()) + to_host = _HostCopyBatch(pin_memory) + # Validate and split the payload before launching any asynchronous copies. + copy_tensor = torch.Tensor.detach if use_async_output else to_host.copy if isinstance(multimodal_outputs_raw, torch.Tensor): # One row per request. The old asserts (`shape[0] == 1` AND # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently @@ -456,7 +520,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "to return multiple tensors per request, use a dict)." ) for i in range(num_reqs): - per_req_payloads.append({"model_outputs": to_host.copy(multimodal_outputs_raw[i])}) + per_req_payloads.append({"model_outputs": copy_tensor(multimodal_outputs_raw[i])}) elif isinstance(multimodal_outputs_raw, list): # One entry per request. The old `len == 1` assert did not check # num_reqs, so a batched step built a length-1 payload list that @@ -469,7 +533,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "to return multiple lists per request, use a dict)." ) for out in multimodal_outputs_raw: - per_req_payloads.append({"model_outputs": to_host.copy(out) if out is not None else None}) + per_req_payloads.append({"model_outputs": copy_tensor(out) if out is not None else None}) elif isinstance(multimodal_outputs_raw, Mapping): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): @@ -481,16 +545,42 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin f"Multimodal output list for key '{key}' has length {len(out)} " f"but expected {num_reqs} (one entry per request)." ) - mm_payload[key] = to_host.copy(out[i]) + mm_payload[key] = copy_tensor(out[i]) elif isinstance(out, torch.Tensor): - mm_payload[key] = to_host.copy(out) + mm_payload[key] = copy_tensor(out) else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") per_req_payloads.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") - # Sync mode hands these payloads out as finished host tensors. - to_host.wait() + snapshot = None + if use_async_output: + # Models may opt in only when all CUDA payload storage is immutable + # across subsequent forwards. Outer model graph capture can still reuse + # these allocations, so require eager outer execution (internal graphs + # that return owned copies are safe). Otherwise keep the snapshot. + copy_payload = ( + _copy_owned_generation_payload_to_cpu_async + if ( + getattr(getattr(self, "model", None), "owns_generation_output_storage", False) is True + and getattr(getattr(self, "model_config", None), "enforce_eager", False) is True + ) + else _snapshot_tensor_payload_to_cpu_async + ) + # Dense waveforms can use one host allocation and one D2H transfer. + # Split only after copying; the CPU rows retain the batch storage. + if isinstance(multimodal_outputs_raw, torch.Tensor) and multimodal_outputs_raw.is_contiguous(): + snapshot = copy_payload( + multimodal_outputs_raw, copy_stream=self.async_output_copy_stream, pin_memory=pin_memory + ) + per_req_payloads = [{"model_outputs": row} for row in snapshot.payload.unbind(0)] + else: + snapshot = copy_payload( + per_req_payloads, copy_stream=self.async_output_copy_stream, pin_memory=pin_memory + ) + per_req_payloads = snapshot.payload + else: + to_host.wait() inter_stage_outputs: list[dict[str, object] | None] | None multimodal_outputs: list[dict[str, object] | None] | None @@ -505,7 +595,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() - if inter_stage_outputs and self._should_accumulate_full_payload_output(): + if inter_stage_outputs and accumulate_full_payload: for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) if req_state is not None and inter_stage_outputs[i]: @@ -528,6 +618,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin ) output.omni_connector_output = self.get_omni_connector_output() + if snapshot is not None: + return _AsyncGenerationOutput(output, snapshot) + if not self.use_async_scheduling: return output