[Core] Defer generation payload host waits to output consumption (#8540)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

.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
 

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)

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])

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).

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

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,
         *,

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)
 

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