Enable Diffusion Image models asynchronous output (#4981)

Signed-off-by: zzh <943967662@qq.com>
Co-authored-by: dengyunyang <584797741@qq.com>

docs/design/feature/async_diffusion_output.md

@@ -0,0 +1,174 @@
+# Async Diffusion Output
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Model Coverage](#model-coverage)
+5. [Configuration](#configuration)
+6. [Related Files](#related-files)
+
+## Overview
+
+The async diffusion output feature moves the D2H (Device→Host) copy and SHM packing from the Worker main thread to a background daemon thread, so the GPU default stream can immediately start the next request's forward pass. This eliminates the GPU bubble caused by synchronous D2H/packing.
+
+**Core benefit**: the previous request's output D2H/packing overlaps with the next request's forward pass.
+
+**Automatically enabled** when `step_execution=False` (default, request-mode). No extra configuration needed. When `step_execution=True` (step-mode), the original synchronous path is used — pump thread and Worker background thread are not started.
+
+## Performance
+
+### HunyuanImage-3.0 (TP4)
+
+| Resolution | Async Off (QPS) | Async On (QPS) | Change |
+|---|---|---|---|
+| 1024×1024 | 0.4773 | 0.4802 | +0.60% |
+| 768×768 | 0.8370 | 0.8533 | +1.95% |
+
+**Analysis**: smaller resolution → larger QPS gain. D2H/packing is relatively fixed overhead; forward is variable (scales with resolution). Smaller images → shorter forward → D2H/packing occupies a larger fraction of total time → more benefit from overlapping it.
+
+## Architecture
+
+### Execution Timeline
+
+```
+Before (synchronous D2H):
+[forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2]
+                                 ^^^^^^^^^^^^^^^^ GPU bubble
+
+After (async D2H):
+[forward req1] [forward req2]
+[D2H+SHM pack req1 (side stream)] [D2H+SHM pack req2 (side stream)]
+                                 ^^^^^^^^^^^^^^^^ bubble eliminated
+```
+
+### Data Flow
+
+```
+                    ┌─────────────────────────────────────────────────────┐
+                    │                 Worker Process                      │
+                    │                                                    │
+  execute_model ──> │ WorkerBusyLoop ──> forward() ──> DiffusionOutput   │
+                    │       │                              │ (GPU tensor)│
+                    │       │                    ┌─────────┴──────────┐  │
+                    │       │                    │  async output path │  │
+                    │       │                    │                    │  │
+                    │       │              compute_done ──> result_mq │  │
+                    │       │                    │                    │  │
+                    │       │                    ▼                    │  │
+                    │       │            AsyncOutputThread            │  │
+                    │       │         (side CUDA stream)             │  │
+                    │       │         wait_event(gpu_event)          │  │
+                    │       │         pack_diffusion_output_shm()    │  │
+                    │       │         D2H + SHM write                │  │
+                    │       │                    │                    │  │
+                    │       │              output_ready ──> result_mq │  │
+                    │       │                    │                    │  │
+                    │       ▼                    ▼                    │  │
+                    │  (dequeue next request)                        │  │
+                    └─────────────────────────────────────────────────────┘
+                                         │
+                                         ▼
+                    ┌─────────────────────────────────────────────────────┐
+                    │              Executor (main process)               │
+                    │                                                    │
+                    │  ResultPumpThread (sole reader of result_mq)       │
+                    │       │                                            │
+                    │       ├── AsyncDiffusionOutput(COMPUTE_DONE)       │
+                    │       │   → resolve _rpc_futures[rpc_id]           │
+                    │       │   → RunnerOutput(async_output_id=...)      │
+                    │       │                                            │
+                    │       ├── AsyncDiffusionOutput(OUTPUT_READY)       │
+                    │       │   → resolve _output_futures[async_output_id]│
+                    │       │   → or batch split via _batch_split_map   │
+                    │       │                                            │
+                    │       └── Non-async message                        │
+                    │           → _sync_result_buffer (for other RPCs)   │
+                    │                                                    │
+                    │  wait_output_ready(async_output_id) → Future       │
+                    └─────────────────────────────────────────────────────┘
+```
+
+### Key Components
+
+1. **AsyncDiffusionOutput** (`data.py`): Protocol envelope for `result_mq`. `kind` field routes messages:
+   - `COMPUTE_DONE` — forward finished, GPU can start next request
+   - `OUTPUT_READY` — D2H/SHM packing finished, final output available
+   - `RPC_RESULT` — ordinary RPC return (including error propagation)
+
+2. **Worker AsyncOutputThread** (`diffusion_worker.py`): Background daemon thread that:
+   - Waits for `gpu_event` (cross-stream sync ensuring forward wrote the tensor)
+   - Runs `pack_diffusion_output_shm()` on a side CUDA stream
+   - Enqueues `OUTPUT_READY` to `result_mq` when packing completes
+
+3. **ResultPumpThread** (`multiproc_executor.py`): Sole reader of `result_mq` in request-mode. Dispatches messages:
+   - `COMPUTE_DONE` / `RPC_RESULT` → resolves `_rpc_futures[rpc_id]`
+   - `OUTPUT_READY` → resolves `_output_futures[async_output_id]` (with batch split via `_batch_split_map`)
+   - Non-async messages → `_sync_result_buffer` for other RPCs
+
+4. **collective_rpc() two-path dispatch** (`multiproc_executor.py`):
+   - **Path 1**: `execute_model` / `execute_model_batch` → generates `rpc_id`, registers Future, waits for pump to deliver `compute_done`
+   - **Path 2**: All other RPCs → `_sync_result_buffer` (pump-fed) or `_result_mq` (step-mode, no pump)
+
+5. **record_device_event** (`platforms/`): Cross-platform GPU event recording for side-stream synchronization. Implemented for CUDA, NPU, ROCm, XPU, MUSA; base class returns `None` (safe no-op).
+
+6. **IPC side-stream D2H** (`ipc.py`): `pack_diffusion_output_shm()` and helpers accept `d2h_stream` parameter; use `pin_memory` + `copy_(non_blocking=True)` on the side stream instead of synchronous `.cpu()`.
+
+### Batch Split
+
+When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
+
+## Model Coverage
+
+### Currently Supported (request-mode, `step_execution=False`)
+
+All models running in request-mode (`step_execution=False`, the default) automatically use async output.
+
+**Verified models**:
+
+| Model | Type | `supports_request_batch` |
+|---|---|---|
+| **HunyuanImage-3.0** | Image | `False` |
+| **Qwen-Image** | Image | `True` |
+
+Other models with `step_execution=False` are also supported but not yet verified.
+
+### Not Yet Supported (step-mode, `step_execution=True`)
+
+When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable.
+
+**Verified**: Helios (`step_execution=True`, does not benefit from this feature).
+
+Other models with `step_execution=True` are also not applicable.
+
+Adapting step-mode requires additional design because:
+
+- Each step's `RunnerOutput` must be available synchronously for the scheduler to advance
+- Intermediate latent tensors may not need D2H (only the final step needs D2H + postprocess)
+- Steps are serial (step N+1 depends on step N output), so the overlap benefit is limited to "final step's D2H overlaps with the next request's forward"
+
+**Planned approach**: first cover only the final step's async D2H, keep intermediate steps synchronous. This matches the request-mode benefit pattern.
+
+## Configuration
+
+No configuration needed. Async output is automatically enabled when `step_execution=False` (default).
+
+| `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread |
+|---|---|---|---|---|
+| `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started |
+| `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started |
+
+## Related Files
+
+- `vllm_omni/diffusion/data.py`: `AsyncDiffusionOutput`, `AsyncOutputKind`, `DiffusionOutput.async_output_id`
+- `vllm_omni/diffusion/worker/diffusion_worker.py`: `WorkerProc._return_result()`, `_async_output_loop()`, `_generate_async_output_id()`
+- `vllm_omni/diffusion/executor/multiproc_executor.py`: `ResultPumpThread`, `collective_rpc()` two-path dispatch, `wait_output_ready()`, `_batch_split_map`
+- `vllm_omni/diffusion/diffusion_engine.py`: `step()` / `step_streaming()` / `add_req_and_wait_for_response()` async output waiting
+- `vllm_omni/diffusion/sched/request_scheduler.py`: `update_from_output()` — `async_output_id` → `FINISHED_COMPLETED`
+- `vllm_omni/diffusion/ipc.py`: `pack_diffusion_output_shm()` side-stream D2H path
+- `vllm_omni/platforms/interface.py`: `OmniPlatform.record_device_event()` base class
+- `vllm_omni/platforms/cuda/platform.py`: `CudaOmniPlatform.record_device_event()`
+- `vllm_omni/platforms/npu/platform.py`: `NPUOmniPlatform.record_device_event()`
+- `vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py`: `get_hunyuan_image3_post_process_func()`, postprocess removal
+- `vllm_omni/diffusion/registry.py`: `_DIFFUSION_POST_PROCESS_FUNCS` postprocess registry

tests/diffusion/test_async_output_worker.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for WorkerProc async output logic."""
+
+import queue
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.data import (
+    AsyncDiffusionOutput,
+    AsyncOutputKind,
+    DiffusionOutput,
+    OmniACK,
+)
+from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
+from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_worker_proc(step_execution=False):
+    """Create a minimal WorkerProc-like object with async output state."""
+    proc = object.__new__(WorkerProc)
+    proc.od_config = MagicMock()
+    proc.od_config.step_execution = step_execution
+    proc.gpu_id = 0
+    proc.result_mq = MagicMock()
+    proc._running = True
+    proc._async_output_queue = queue.Queue() if not step_execution else None
+    proc._async_output_thread = None
+    return proc
+
+
+class TestGenerateAsyncOutputId:
+    """Test _generate_async_output_id uniqueness."""
+
+    def test_returns_hex_string(self):
+        id1 = WorkerProc._generate_async_output_id()
+        assert isinstance(id1, str)
+        assert len(id1) == 32  # uuid4 hex is 32 chars
+
+    def test_ids_are_unique(self):
+        ids = {WorkerProc._generate_async_output_id() for _ in range(100)}
+        assert len(ids) == 100
+
+
+class TestReturnResultAsyncPath:
+    """Test _return_result in request-mode (step_execution=False)."""
+
+    def test_omni_ack_bypasses_async_path(self):
+        """OmniACK should always use sync path regardless of step_execution."""
+        proc = _make_worker_proc(step_execution=False)
+        ack = OmniACK(task_id="task1", status="done")
+
+        proc._return_result(ack)
+        proc.result_mq.enqueue.assert_called_once_with(ack)
+        assert proc._async_output_queue.empty()  # nothing enqueued for bg thread
+
+    def test_diffusion_output_triggers_async_path(self, mocker):
+        """DiffusionOutput in request-mode should enqueue compute_done + bg work."""
+        proc = _make_worker_proc(step_execution=False)
+
+        mock_platform = mocker.MagicMock()
+        mock_platform.record_device_event.return_value = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+            mock_platform,
+        )
+
+        output = DiffusionOutput(output="data")
+        proc._return_result(output, rpc_id="1")
+
+        # Should have enqueued COMPUTE_DONE to result_mq
+        enqueued_msg = proc.result_mq.enqueue.call_args[0][0]
+        assert isinstance(enqueued_msg, AsyncDiffusionOutput)
+        assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE
+        assert enqueued_msg.rpc_id == "1"
+        assert enqueued_msg.async_output_id is not None
+
+        # Should have placed work item in async_output_queue
+        assert not proc._async_output_queue.empty()
+        item = proc._async_output_queue.get_nowait()
+        assert item[0] is output
+        assert item[1] == enqueued_msg.async_output_id
+
+    def test_batch_runner_output_triggers_async_path(self, mocker):
+        """BatchRunnerOutput in request-mode should also go through async path."""
+        proc = _make_worker_proc(step_execution=False)
+
+        mock_platform = mocker.MagicMock()
+        mock_platform.record_device_event.return_value = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+            mock_platform,
+        )
+
+        batch = BatchRunnerOutput(
+            runner_outputs=[RunnerOutput(request_id="r1", finished=True)],
+        )
+        proc._return_result(batch)
+
+        enqueued_msg = proc.result_mq.enqueue.call_args[0][0]
+        assert isinstance(enqueued_msg, AsyncDiffusionOutput)
+        assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE
+
+    def test_none_result_mq_skips_all_paths(self):
+        """If result_mq is None, _return_result should do nothing."""
+        proc = _make_worker_proc(step_execution=False)
+        proc.result_mq = None
+
+        proc._return_result(DiffusionOutput())
+        # No crash, no enqueue attempted
+
+
+class TestReturnResultSyncPath:
+    """Test _return_result in step-execution mode (sync)."""
+
+    def test_diffusion_output_uses_sync_path_when_step_execution(self, mocker):
+        """In step_execution mode, DiffusionOutput goes through sync SHM pack."""
+        proc = _make_worker_proc(step_execution=True)
+        proc._async_output_queue = None  # step mode doesn't have async queue
+
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+        )
+        output = DiffusionOutput(output="data")
+        proc._return_result(output)
+
+        mock_pack.assert_called_once_with(output)
+        proc.result_mq.enqueue.assert_called_once_with(output)
+
+    def test_non_diffusion_output_skips_async_in_request_mode(self, mocker):
+        """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path."""
+        proc = _make_worker_proc(step_execution=False)
+
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+        )
+        # A plain dict is neither DiffusionOutput nor BatchRunnerOutput
+        output = {"type": "other"}
+        proc._return_result(output)
+
+        mock_pack.assert_called_once_with(output)
+        proc.result_mq.enqueue.assert_called_once_with(output)
+
+
+class TestAsyncOutputLoopLogic:
+    """Test _async_output_loop background thread logic (mocked, no device)."""
+
+    def test_bg_thread_enqueues_output_ready_after_pack(self, mocker):
+        """The real _async_output_loop packs via SHM with d2h_stream and
+        enqueues an OUTPUT_READY message."""
+        import threading
+        import time
+
+        proc = _make_worker_proc(step_execution=False)
+
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+        )
+        # Mock torch to avoid device initialization
+        mock_torch = mocker.MagicMock()
+        mock_stream = mocker.MagicMock()
+        mock_torch.Stream.return_value = mock_stream
+        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+        mock_torch.device.return_value = mocker.MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.torch",
+            mock_torch,
+        )
+
+        output = DiffusionOutput(output="data")
+        gpu_event = mocker.MagicMock()
+        proc._async_output_queue.put((output, "abc123", gpu_event))
+
+        # Run the real _async_output_loop in a background thread.
+        t = threading.Thread(target=proc._async_output_loop, daemon=True)
+        t.start()
+
+        # Give the loop time to process, then unblock and stop.
+        time.sleep(0.5)
+        proc._running = False
+        proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
+        t.join(timeout=2.0)
+
+        mock_pack.assert_called()
+        # Verify an OUTPUT_READY with the correct async_output_id was enqueued.
+        output_ready_enqueued = any(
+            c[0][0].kind == AsyncOutputKind.OUTPUT_READY and c[0][0].async_output_id == "abc123"
+            for c in proc.result_mq.enqueue.call_args_list
+        )
+        assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"

tests/diffusion/test_ipc_async.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for IPC async D2H stream packing."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.ipc import (
+    _pack_diffusion_fields,
+    _pack_tensor_if_large,
+    _pack_value_if_large,
+    pack_diffusion_output_shm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestPackTensorIfLargeWithD2hStream:
+    """Test _pack_tensor_if_large with mock d2h_stream parameter."""
+
+    def test_small_tensor_not_packed_with_stream(self, mocker):
+        """Small tensors pass through even when d2h_stream is provided."""
+        small_tensor = torch.zeros(10)  # way below threshold
+        mock_stream = mocker.MagicMock()
+        result = _pack_tensor_if_large(small_tensor, d2h_stream=mock_stream)
+        # Small tensor should be returned as-is (not a SHM handle dict)
+        assert isinstance(result, torch.Tensor)
+
+    def test_large_tensor_packed_without_stream(self, mocker):
+        """Large tensor is packed via SHM when no d2h_stream."""
+        large_tensor = torch.randn(500_000)  # > 1MB for float32
+        result = _pack_tensor_if_large(large_tensor, d2h_stream=None)
+        assert isinstance(result, dict)
+        assert result.get("__tensor_shm__") is True
+
+    def test_pack_value_if_large_passes_stream_to_tensor_pack(self, mocker):
+        """_pack_value_if_large propagates d2h_stream to _pack_tensor_if_large."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        tensor = torch.zeros(1)
+        d2h = mocker.MagicMock()
+        _pack_value_if_large(tensor, d2h_stream=d2h)
+        mock_pack.assert_called_once_with(tensor, d2h_stream=d2h)
+
+    def test_pack_value_if_large_passes_stream_through_dict(self, mocker):
+        """d2h_stream propagates through dict container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = {"image": torch.zeros(1), "video": torch.zeros(1)}
+        _pack_value_if_large(data, d2h_stream=d2h)
+        assert mock_pack.call_count == 2
+        for call in mock_pack.call_args_list:
+            assert call.kwargs.get("d2h_stream") is d2h or call.args[1] is d2h
+
+    def test_pack_value_if_large_passes_stream_through_list(self, mocker):
+        """d2h_stream propagates through list container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = [torch.zeros(1), torch.zeros(1)]
+        _pack_value_if_large(data, d2h_stream=d2h)
+        assert mock_pack.call_count == 2
+
+    def test_pack_value_if_large_passes_stream_through_tuple(self, mocker):
+        """d2h_stream propagates through tuple container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = (torch.zeros(1), torch.zeros(1))
+        result = _pack_value_if_large(data, d2h_stream=d2h)
+        assert isinstance(result, tuple)
+        assert mock_pack.call_count == 2
+
+
+class TestPackDiffusionFieldsWithD2hStream:
+    """Test _pack_diffusion_fields propagates d2h_stream."""
+
+    def test_output_field_packed_with_stream(self, mocker):
+        mock_pack_value = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_value_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput(output=torch.zeros(1))
+        _pack_diffusion_fields(out, d2h_stream=d2h)
+        mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h)
+
+    def test_trajectory_fields_packed_with_stream(self, mocker):
+        mock_pack_tensor = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput(
+            trajectory_latents=torch.zeros(1),
+            trajectory_log_probs=torch.zeros(1),
+        )
+        _pack_diffusion_fields(out, d2h_stream=d2h)
+        # Should pack trajectory_latents and trajectory_log_probs
+        assert mock_pack_tensor.call_count == 2
+
+
+class TestPackDiffusionOutputShmWithD2hStream:
+    """Test pack_diffusion_output_shm top-level dispatch with d2h_stream."""
+
+    def test_diffusion_output_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput()
+        pack_diffusion_output_shm(out, d2h_stream=d2h)
+        mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h)
+
+    def test_rpc_envelope_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        result = DiffusionOutput()
+        envelope = {"type": "diffusion_rpc_result", "result": result}
+        pack_diffusion_output_shm(envelope, d2h_stream=d2h)
+        # envelope["result"] is replaced in-place by pack_diffusion_output_shm;
+        # assert with the original reference, not the replaced one.
+        mock_pack_fields.assert_called_once_with(
+            result,
+            d2h_stream=d2h,
+        )
+
+    def test_runner_output_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+
+        from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+        result = DiffusionOutput()
+        runner_out = RunnerOutput(
+            request_id="req1",
+            result=result,
+        )
+        pack_diffusion_output_shm(runner_out, d2h_stream=d2h)
+        # runner_out.result is replaced in-place by pack_diffusion_output_shm;
+        # assert with the original reference, not the replaced one.
+        mock_pack_fields.assert_called_once_with(
+            result,
+            d2h_stream=d2h,
+        )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -56,7 +56,7 @@ def _make_executor(num_gpus: int = 1):
 
     Returns ``(executor, request_queue, result_queue)``.
     """
-    od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False)
+    od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False, step_execution=True)
     executor = object.__new__(MultiprocDiffusionExecutor)
     executor.od_config = od_cfg
 
@@ -616,6 +616,7 @@ class TestMultiprocExecutorRaisesEngineDeadError:
 
     def test_collective_rpc_raises_when_is_failed(self):
         executor = object.__new__(MultiprocDiffusionExecutor)
+        executor.od_config = SimpleNamespace(step_execution=True)
         executor._closed = False
         executor._broadcast_mq = MagicMock()
         executor._result_mq = MagicMock()
@@ -661,7 +662,7 @@ class TestMultiprocExecutorStepStreamingOutput:
 
     def test_execute_step_allows_streaming_output_mode(self):
         executor, req_q, res_q = _make_executor()
-        executor.od_config = SimpleNamespace(streaming_output=True)  # pyright: ignore[reportAttributeAccessIssue]
+        executor.od_config = SimpleNamespace(streaming_output=True, step_execution=True)  # pyright: ignore[reportAttributeAccessIssue]
         runner_outputs = [
             RunnerOutput(
                 request_id="sched-stream",
@@ -939,6 +940,14 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
         executor._finalizer = lambda: None
+        # ------------------------------------------------------------------
+        # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
+        # over them, so they need to exist even when constructed via __new__.
+        executor._pump_stop = threading.Event()
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._batch_split_map = {}
 
         proc = _make_short_lived_process()
         executor._processes = [proc]

tests/diffusion/test_result_pump.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MultiprocDiffusionExecutor async result pump and wait_output_ready."""
+
+import concurrent.futures
+import queue
+import threading
+import time
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.data import AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_executor(step_execution=False):
+    """Create a minimal MultiprocDiffusionExecutor-like object with pump state."""
+    from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+
+    od_config = MagicMock()
+    od_config.step_execution = step_execution
+
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = od_config
+    executor._rpc_id_counter = 0
+    executor._rpc_id_lock = threading.Lock()
+    executor._rpc_futures = {}
+    executor._output_futures = {}
+    executor._completed_outputs = {}
+    executor._batch_split_map = {}
+    executor._futures_lock = threading.RLock()
+    executor._pump_running = False
+    executor._pump_stop = threading.Event()
+    executor._sync_result_buffer = queue.Queue()
+    executor._result_mq = MagicMock()
+    executor._broadcast_mq = MagicMock()
+    executor._closed = False
+    executor._is_failed = False
+    executor._finalizer = MagicMock()  # no-op in tests
+    executor._shutdown_cleaner = None
+    executor._processes = []
+    return executor
+
+
+def _feed_one_msg_to_pump(executor, msg):
+    """Run _result_pump in a daemon thread, feed one *msg*, then stop."""
+    call_count = [0]
+
+    def mock_dequeue(timeout=None):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            return msg
+        executor._pump_stop.set()
+        time.sleep(0.05)
+        raise TimeoutError
+
+    executor._result_mq.dequeue = mock_dequeue
+    t = threading.Thread(target=executor._result_pump, daemon=True)
+    t.start()
+    t.join(timeout=2.0)
+
+
+@pytest.fixture(autouse=True)
+def _mock_unpack(mocker):
+    """Real _result_pump calls unpack_diffusion_output_shm; mock it away."""
+    mocker.patch(
+        "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm",
+    )
+
+
+class TestNextRpcId:
+    """Test _next_rpc_id counter."""
+
+    def test_counter_increments(self):
+        executor = _make_executor()
+        id1 = executor._next_rpc_id()
+        id2 = executor._next_rpc_id()
+        id3 = executor._next_rpc_id()
+        assert id1 == "1"
+        assert id2 == "2"
+        assert id3 == "3"
+
+    def test_counter_is_threadsafe(self):
+        executor = _make_executor()
+        ids = []
+
+        def get_id():
+            ids.append(executor._next_rpc_id())
+
+        threads = [threading.Thread(target=get_id) for _ in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # All IDs should be unique
+        assert len(set(ids)) == 10
+
+
+class TestWaitOutputReady:
+    """Test wait_output_ready future creation and caching."""
+
+    def test_returns_new_future_when_not_cached(self):
+        executor = _make_executor()
+        fut = executor.wait_output_ready("abc123")
+        assert isinstance(fut, concurrent.futures.Future)
+        assert not fut.done()
+
+    def test_future_resolves_when_output_arrives(self):
+        executor = _make_executor()
+        fut = executor.wait_output_ready("abc123")
+        output = DiffusionOutput(output="data")
+
+        # Simulate pump resolving the future
+        with executor._futures_lock:
+            executor._output_futures.pop("abc123")
+        if not fut.done():
+            fut.set_result(output)
+
+        assert fut.result(timeout=1.0) is output
+
+    def test_returns_cached_future_when_already_completed(self):
+        executor = _make_executor()
+        output = DiffusionOutput(output="cached_data")
+        fut = concurrent.futures.Future()
+        fut.set_result(output)
+        with executor._futures_lock:
+            executor._completed_outputs["abc123"] = fut
+
+        fut = executor.wait_output_ready("abc123")
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+
+    def test_removes_from_cache_after_retrieval(self):
+        executor = _make_executor()
+        output = DiffusionOutput(output="cached_data")
+        with executor._futures_lock:
+            executor._completed_outputs["abc123"] = output
+
+        executor.wait_output_ready("abc123")
+        # Second call should not find cached result
+        with executor._futures_lock:
+            assert "abc123" not in executor._completed_outputs
+
+
+class TestResultPumpDispatch:
+    """Test _result_pump message routing (running the real pump in a thread)."""
+
+    def test_non_async_message_placed_in_sync_buffer(self):
+        executor = _make_executor()
+        msg = DiffusionOutput(output="sync_result")
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert not executor._sync_result_buffer.empty()
+        retrieved = executor._sync_result_buffer.get_nowait()
+        assert isinstance(retrieved, DiffusionOutput)
+
+    def test_compute_done_routes_to_rpc_future(self):
+        executor = _make_executor()
+        rpc_id = "42"
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._rpc_futures[rpc_id] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.COMPUTE_DONE,
+            rpc_id=rpc_id,
+            async_output_id="abc",
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        result = fut.result(timeout=1.0)
+        assert result.kind == AsyncOutputKind.COMPUTE_DONE
+
+    def test_output_ready_routes_to_output_future(self):
+        executor = _make_executor()
+        async_output_id = "abc123"
+        output = DiffusionOutput(output="final")
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures[async_output_id] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            output=output,
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+
+    def test_output_ready_with_error_routes_to_future_as_exception(self):
+        executor = _make_executor()
+        async_output_id = "abc123"
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures[async_output_id] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            error="Background D2H/SHM packing failed",
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        with pytest.raises(RuntimeError, match="Background D2H/SHM packing failed"):
+            fut.result(timeout=1.0)
+
+    def test_output_ready_caches_when_no_future_waiting(self):
+        """When OUTPUT_READY arrives but no future is waiting, result is cached."""
+        executor = _make_executor()
+        async_output_id = "abc123"
+        output = DiffusionOutput(output="orphan")
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            output=output,
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        # Later call to wait_output_ready should find it cached
+        fut = executor.wait_output_ready(async_output_id)
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+
+
+class TestShutdownCleansUpFutures:
+    """Test that shutdown cancels pending futures."""
+
+    def test_shutdown_sets_exception_on_pending_futures(self):
+        executor = _make_executor()
+
+        rpc_fut = concurrent.futures.Future()
+        output_fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._rpc_futures["1"] = rpc_fut
+            executor._output_futures["abc"] = output_fut
+
+        executor.shutdown()
+
+        assert rpc_fut.done()
+        with pytest.raises(RuntimeError, match="Executor shut down"):
+            rpc_fut.result(timeout=1.0)
+
+        assert output_fut.done()
+        with pytest.raises(RuntimeError, match="Executor shut down"):
+            output_fut.result(timeout=1.0)
+
+        assert len(executor._rpc_futures) == 0
+        assert len(executor._output_futures) == 0

tests/platforms/test_record_device_event.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OmniPlatform.record_device_event implementations."""
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniPlatformRecordDeviceEventInterface:
+    """Test that the base OmniPlatform returns None (safe no-op fallback)."""
+
+    def test_base_class_returns_none(self):
+        from vllm_omni.platforms.interface import OmniPlatform
+
+        assert OmniPlatform.record_device_event() is None
+
+
+class TestCudaOmniPlatformRecordDeviceEvent:
+    """Test CudaOmniPlatform.record_device_event with mocked torch.Event."""
+
+    def test_records_event_successfully(self, mocker):
+        from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+        mock_event = mocker.MagicMock()
+        mocker.patch("torch.Event", return_value=mock_event)
+
+        result = CudaOmniPlatform.record_device_event()
+        assert result is mock_event
+        mock_event.record.assert_called_once()
+
+    def test_returns_none_on_failure(self, mocker):
+        from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+        mocker.patch("torch.Event", side_effect=RuntimeError("no CUDA"))
+
+        result = CudaOmniPlatform.record_device_event()
+        assert result is None
+
+
+class TestNPUOmniPlatformRecordDeviceEvent:
+    """Test NPUOmniPlatform.record_device_event with mocked torch.npu."""
+
+    def test_returns_none_on_failure(self, mocker):
+        """When torch.npu.current_stream().synchronize() fails, return None."""
+        try:
+            from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+        except ModuleNotFoundError:
+            pytest.skip("vllm_ascend not available")
+
+        mock_torch = mocker.MagicMock()
+        mock_torch.npu.current_stream.return_value.synchronize.side_effect = RuntimeError("no NPU")
+        mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch)
+
+        result = NPUOmniPlatform.record_device_event()
+        assert result is None
+
+    def test_synchronizes_stream_then_records_event(self, mocker):
+        """NPU should sync stream first, then create and record an Event."""
+        try:
+            from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+        except ModuleNotFoundError:
+            pytest.skip("vllm_ascend not available")
+
+        mock_event = mocker.MagicMock()
+        mock_stream = mocker.MagicMock()
+        mock_torch = mocker.MagicMock()
+        mock_torch.npu.current_stream.return_value = mock_stream
+        mock_torch.npu.Event.return_value = mock_event
+        mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch)
+
+        result = NPUOmniPlatform.record_device_event()
+
+        # Stream should be synced first (HCCL ordering)
+        mock_stream.synchronize.assert_called_once()
+        # Then event should be created and recorded
+        mock_torch.npu.Event.assert_called_once()
+        mock_event.record.assert_called_once()
+        assert result is mock_event

vllm_omni/diffusion/data.py

@@ -1273,6 +1273,7 @@ class DiffusionOutput:
     trajectory_latents: torch.Tensor | dict[str, Any] | None = None
     trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
     trajectory_decoded: list[Image.Image] | None = None
+    async_output_id: str | None = None
     error: str | None = None
     error_status_code: int | None = None
     error_type: str | None = None
@@ -1333,6 +1334,37 @@ class DiffusionOutput:
         )
 
 
+class AsyncOutputKind(Enum):
+    """Message kind for ``AsyncDiffusionOutput`` — routes async results.
+
+    * ``RPC_RESULT`` — ordinary RPC return (sleep, wake, profile, etc.)
+    * ``COMPUTE_DONE`` — worker forward finished, GPU can start next request
+    * ``OUTPUT_READY`` — background D2H/SHM packing finished, final output
+      is available via ``async_output_id``
+    """
+
+    RPC_RESULT = "rpc_result"
+    COMPUTE_DONE = "compute_done"
+    OUTPUT_READY = "output_ready"
+
+
+@dataclass
+class AsyncDiffusionOutput:
+    """Async protocol envelope for ``result_mq`` messages.
+
+    In request-mode (``step_execution=False``), all ``result_mq``
+    messages use this envelope.  The ``kind`` field routes the message
+    to the correct consumer.
+    """
+
+    kind: AsyncOutputKind
+    rpc_id: str | None = None
+    async_output_id: str | None = None
+    result: Any | None = None
+    output: DiffusionOutput | None = None
+    error: str | None = None
+
+
 class DiffusionRequestAbortedError(RuntimeError):
     """Raised when a diffusion request ends via user-visible abort."""
 

vllm_omni/diffusion/diffusion_engine.py

@@ -55,6 +55,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+_ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
+
 __all__ = [
     "DiffusionEngine",
     "DiffusionExecutionMode",
@@ -288,6 +290,10 @@ class DiffusionEngine:
         generator = self.async_add_req_and_stream_response(request)
         async for output in generator:
             exec_total_time = time.perf_counter() - exec_start_time
+            # Async mode: wait for background D2H/SHM to complete.
+            if output.async_output_id:
+                fut = self.executor.wait_output_ready(output.async_output_id)
+                output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT)
             postprocess_start_time = time.perf_counter()
             formatted_outputs = self.postprocess_output(request, output)
             postprocess_time = time.perf_counter() - postprocess_start_time
@@ -744,11 +750,15 @@ class DiffusionEngine:
                     raise ValueError("Sync func should receive one result at one time")
                 if target_request_id in finished_req_ids:
                     req_output = runner_output.get_request_output(target_request_id)
-                    return self._finalize_finished_request(
+                    output = self._finalize_finished_request(
                         target_request_id,
                         runner_output=req_output,
                         missing_result_error="Diffusion execution finished without a final output.",
                     )
+                    if output.async_output_id:
+                        fut = self.executor.wait_output_ready(output.async_output_id)
+                        output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT)
+                    return output
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop profiling on all diffusion workers.
@@ -1036,4 +1046,7 @@ class DiffusionEngine:
         if runner_output is not None and runner_output.result is not None:
             return runner_output.result
 
+        if runner_output is not None and runner_output.async_output_id is not None:
+            return DiffusionOutput(async_output_id=runner_output.async_output_id)
+
         return DiffusionOutput(error=missing_result_error)

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import concurrent.futures
 import multiprocessing as mp
 import multiprocessing.connection
+import queue
 import threading
 import time
 import weakref
@@ -15,7 +17,7 @@ from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQ
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
 from vllm_omni.diffusion.worker import WorkerProc
@@ -93,6 +95,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             self.shutdown()
             raise
 
+        # Async output: result pump thread for dispatching AsyncDiffusionOutput
+        # messages in request-mode (step_execution=False).
+        self._rpc_id_counter = 0
+        self._rpc_id_lock = threading.Lock()
+        self._rpc_futures: dict[str, concurrent.futures.Future[AsyncDiffusionOutput]] = {}
+        self._output_futures: dict[str, concurrent.futures.Future[DiffusionOutput]] = {}
+        self._completed_outputs: dict[str, concurrent.futures.Future[DiffusionOutput]] = {}
+        self._batch_split_map: dict[str, dict[str, str]] = {}  # batch_id -> {per_req_id: request_id}
+        self._futures_lock = threading.RLock()
+        self._pump_running = False
+        self._pump_stop = threading.Event()
+        # When pump is active it is the sole reader of result_mq; non-async
+        # messages are placed here for collective_rpc() to consume.
+        self._sync_result_buffer: queue.Queue = queue.Queue()
+        if not self.od_config.step_execution:
+            self._start_result_pump()
+
         self._start_worker_monitor()
 
     def _init_broadcast_queue(self, num_workers: int) -> MessageQueue:
@@ -116,7 +135,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise RuntimeError("Broadcast queue is closed")
 
     def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any:
-        """Block until one result message, polling ``_is_failed`` between chunk timeouts."""
+        """Block until one result message, polling ``_is_failed`` between chunk timeouts.
+
+        When async output is enabled, pump is the sole reader of result_mq;
+        non-async messages are placed in _sync_result_buffer, so this method
+        reads from there instead of result_mq directly.
+        """
         while True:
             if deadline is None:
                 chunk_timeout = _DEQUEUE_TIMEOUT_S
@@ -125,12 +149,20 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 if remaining <= 0:
                     raise TimeoutError(f"RPC call to {method} timed out.")
                 chunk_timeout = min(_DEQUEUE_TIMEOUT_S, remaining)
-            try:
-                return self._result_mq.dequeue(timeout=chunk_timeout)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
-            except (TimeoutError, zmq.error.Again):
-                if self._is_failed:
-                    raise EngineDeadError()
-                continue
+            if not self.od_config.step_execution:
+                try:
+                    return self._sync_result_buffer.get(timeout=chunk_timeout)
+                except queue.Empty:
+                    if self._is_failed or self._closed:
+                        raise EngineDeadError()
+                    continue
+            else:
+                try:
+                    return self._result_mq.dequeue(timeout=chunk_timeout)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
+                except (TimeoutError, zmq.error.Again):
+                    if self._is_failed:
+                        raise EngineDeadError()
+                    continue
 
     @staticmethod
     def _raise_for_rpc_error_dict(response: Any) -> None:
@@ -318,6 +350,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         """Adapt request-mode scheduler output to worker execute_model RPCs.
 
         Returns a BatchRunnerOutput with one RunnerOutput per scheduled request.
+        In async mode the result may carry async_output_id instead of the final output.
         """
         from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 
@@ -333,16 +366,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     unique_reply_rank=0,
                     exec_all_ranks=True,
                 )
-                if not isinstance(result, DiffusionOutput):
-                    raise RuntimeError(f"Unexpected response type: {type(result)!r}")
-                runner_outputs.append(
-                    RunnerOutput(
-                        request_id=new_req.request_id,
-                        step_index=None,
-                        finished=True,
-                        result=result,
+                if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE:
+                    runner_outputs.append(
+                        RunnerOutput(
+                            request_id=new_req.request_id,
+                            step_index=None,
+                            finished=True,
+                            result=None,
+                            async_output_id=result.async_output_id,
+                        )
                     )
-                )
+                elif isinstance(result, DiffusionOutput):
+                    runner_outputs.append(
+                        RunnerOutput(
+                            request_id=new_req.request_id,
+                            step_index=None,
+                            finished=True,
+                            result=result,
+                        )
+                    )
+                else:
+                    raise RuntimeError(f"Unexpected response type: {type(result)!r}")
             except Exception as exc:
                 runner_outputs.append(
                     RunnerOutput(
@@ -361,8 +405,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         A scheduler wave with one request is the conservative serial case and
         uses the single-request worker RPC. Waves with multiple requests use the
         fused request-batch RPC and require pipeline request-batch support.
+
+        When async output is enabled, per-request async_output_ids are
+        propagated through RunnerOutput so the engine waits in
+        step_streaming() instead of blocking the busy loop here.
         """
-        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
+        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 
         self._ensure_open()
         if len(scheduler_output.scheduled_new_reqs) <= 1:
@@ -374,6 +422,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             unique_reply_rank=0,
             exec_all_ranks=True,
         )
+        if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE:
+            # Propagate async_output_id to per-request RunnerOutputs so the
+            # engine waits in step_streaming() instead of blocking here.
+            batch_id = result.async_output_id
+            per_req_map: dict[str, str] = {}
+            runner_outputs: list[RunnerOutput] = []
+            for new_req in scheduler_output.scheduled_new_reqs:
+                per_req_id = f"{batch_id}/{new_req.request_id}"
+                per_req_map[per_req_id] = new_req.request_id
+                runner_outputs.append(
+                    RunnerOutput(
+                        request_id=new_req.request_id,
+                        step_index=None,
+                        finished=True,
+                        result=None,
+                        async_output_id=per_req_id,
+                    )
+                )
+            with self._futures_lock:
+                self._batch_split_map[batch_id] = per_req_map
+            return BatchRunnerOutput.from_list(runner_outputs)
         if not isinstance(result, BatchRunnerOutput):
             raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}")
         return result
@@ -408,11 +477,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         deadline = None if timeout is None else time.monotonic() + timeout
         kwargs = kwargs or {}
 
-        # Prepare RPC request message. When unique_reply_rank is None, all
-        # workers must execute the RPC but only rank 0 can reply (it's the
-        # only one with a result_mq). Collect detailed rank statuses only for
-        # this control-plane all-rank path; forward-path exec_all_ranks RPCs
-        # avoid the per-step host object gather.
         execute_all_ranks = unique_reply_rank is None or exec_all_ranks
         collect_rank_status = unique_reply_rank is None
         rpc_request = {
@@ -425,11 +489,32 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             "collect_rank_status": collect_rank_status,
         }
 
+        # ── Path 1: async execute_model / execute_model_batch ──
+        # Request-mode methods use the compute_done/output_ready split
+        # so the D2H copy runs on a side stream in the worker background
+        # thread while the default stream is free for the next forward.
+        # pump routes the result back to a Future via rpc_id.
+        if not self.od_config.step_execution and method in ("execute_model", "execute_model_batch"):
+            rpc_id = self._next_rpc_id()
+            rpc_request["rpc_id"] = rpc_id
+            fut: concurrent.futures.Future = concurrent.futures.Future()
+            with self._futures_lock:
+                self._rpc_futures[rpc_id] = fut
+            self._broadcast_mq.enqueue(rpc_request)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
+            try:
+                return fut.result(timeout=timeout)
+            except concurrent.futures.TimeoutError:
+                with self._futures_lock:
+                    self._rpc_futures.pop(rpc_id, None)
+                raise TimeoutError(f"RPC call to {method} timed out.")
+
+        # ── Path 2: step-execution or other non-execute_model RPCs ──
+        # _dequeue_one_with_failure_polling reads from _sync_result_buffer
+        # when pump is active (request-mode), or result_mq directly otherwise.
         try:
             # Broadcast RPC request to all workers via unified message queue
             self._broadcast_mq.enqueue(rpc_request)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
 
-            # Only rank 0 has a result_mq, so we always expect exactly 1 response
             num_responses = 1
 
             responses = []
@@ -450,6 +535,123 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             logger.error(f"RPC call failed: {e}")
             raise
 
+    # ------------------------------------------------------------------
+    # Async output: result pump
+    # ------------------------------------------------------------------
+
+    def _start_result_pump(self) -> None:
+        self._pump_running = True
+        self._pump_stop.clear()
+        self._result_pump_thread = threading.Thread(target=self._result_pump, daemon=True, name="DiffusionResultPump")
+        self._result_pump_thread.start()
+        logger.info("Async result pump started")
+
+    def _result_pump(self) -> None:
+        """Sole reader of result_mq when async output is enabled.
+
+        Dispatches AsyncDiffusionOutput messages to the appropriate future:
+        * RPC_RESULT / COMPUTE_DONE → _rpc_futures[rpc_id]
+        * OUTPUT_READY → _output_futures[async_output_id]
+        """
+        while not self._pump_stop.is_set():
+            try:
+                msg = self._result_mq.dequeue(timeout=1.0)
+            except TimeoutError:
+                if self._is_failed:
+                    break
+                continue
+            except Exception:
+                logger.exception("Result pump dequeue failed")
+                if self._is_failed:
+                    break
+                continue
+
+            if not isinstance(msg, AsyncDiffusionOutput):
+                # Non-async message: place into the sync buffer for
+                # collective_rpc() to consume via Path 2.
+                self._sync_result_buffer.put(msg)
+                continue
+
+            if msg.kind in (AsyncOutputKind.RPC_RESULT, AsyncOutputKind.COMPUTE_DONE):
+                with self._futures_lock:
+                    fut = self._rpc_futures.pop(msg.rpc_id, None) if msg.rpc_id else None
+                if fut is not None and not fut.done():
+                    if msg.error:
+                        fut.set_exception(RuntimeError(msg.error))
+                    else:
+                        fut.set_result(msg)
+            elif msg.kind == AsyncOutputKind.OUTPUT_READY:
+                batch_id = msg.async_output_id
+                with self._futures_lock:
+                    per_req_map = self._batch_split_map.pop(batch_id, None) if batch_id else None
+                if per_req_map is not None:
+                    # Batch result: split into per-request DiffusionOutputs.
+                    try:
+                        unpack_diffusion_output_shm(msg.output)
+                    except Exception:
+                        logger.exception("SHM unpack failed for batch %s", batch_id)
+                    batch_output = msg.output
+                    for per_req_id, req_id in per_req_map.items():
+                        req_output = batch_output.get_request_output(req_id)
+                        per_req_result: DiffusionOutput
+                        if req_output is not None and req_output.result is not None:
+                            per_req_result = req_output.result
+                        elif msg.error:
+                            per_req_result = DiffusionOutput(error=msg.error)
+                        else:
+                            per_req_result = DiffusionOutput(error="No output result for batch request")
+                        fut: concurrent.futures.Future = concurrent.futures.Future()
+                        fut.set_result(per_req_result)
+                        with self._futures_lock:
+                            pending = self._output_futures.pop(per_req_id, None)
+                            if pending is not None and not pending.done():
+                                pending.set_result(per_req_result)
+                            else:
+                                self._completed_outputs[per_req_id] = fut
+                else:
+                    with self._futures_lock:
+                        fut = self._output_futures.pop(batch_id, None) if batch_id else None
+                    if fut is not None and not fut.done():
+                        if msg.error:
+                            fut.set_exception(RuntimeError(msg.error))
+                        else:
+                            try:
+                                unpack_diffusion_output_shm(msg.output)
+                            except Exception as e:
+                                logger.exception("SHM unpack failed in result pump")
+                                fut.set_exception(e)
+                                continue
+                            fut.set_result(msg.output)
+                    elif batch_id:
+                        fut = concurrent.futures.Future()
+                        if msg.error:
+                            fut.set_exception(RuntimeError(msg.error))
+                        else:
+                            try:
+                                unpack_diffusion_output_shm(msg.output)
+                            except Exception as e:
+                                logger.exception("SHM unpack failed in result pump (cached)")
+                                fut.set_exception(e)
+                            else:
+                                fut.set_result(msg.output)
+                        with self._futures_lock:
+                            self._completed_outputs[batch_id] = fut
+
+    def _next_rpc_id(self) -> str:
+        with self._rpc_id_lock:
+            self._rpc_id_counter += 1
+            return str(self._rpc_id_counter)
+
+    def wait_output_ready(self, async_output_id: str) -> concurrent.futures.Future[DiffusionOutput]:
+        """Return a Future that resolves when the async output is ready."""
+        with self._futures_lock:
+            cached = self._completed_outputs.pop(async_output_id, None)
+            if cached is not None:
+                return cached
+            fut: concurrent.futures.Future = concurrent.futures.Future()
+            self._output_futures[async_output_id] = fut
+        return fut
+
     def check_health(self) -> None:
         if self._is_failed:
             raise EngineDeadError()
@@ -461,10 +663,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
     def shutdown(self) -> None:
         self._closed = True
+        self._pump_stop.set()
         try:
             self._finalizer()
         finally:
             self._broadcast_mq = None
             self._result_mq = None
+            with self._futures_lock:
+                for fut in self._rpc_futures.values():
+                    if not fut.done():
+                        fut.set_exception(RuntimeError("Executor shut down"))
+                for fut in self._output_futures.values():
+                    if not fut.done():
+                        fut.set_exception(RuntimeError("Executor shut down"))
+                self._rpc_futures.clear()
+                self._output_futures.clear()
+                self._batch_split_map.clear()
             self._shutdown_cleaner = None
             self._processes = []

vllm_omni/diffusion/ipc.py

@@ -21,23 +21,43 @@ _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
 DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
 
 
-def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
+def _tensor_to_shm(
+    tensor: torch.Tensor,
+    d2h_stream: torch.Stream | None = None,
+) -> dict[str, Any]:
     """Copy a tensor into POSIX shared memory and return a metadata handle.
 
     The shared memory segment remains alive after this call (the local fd is
     closed, but the segment persists until ``_tensor_from_shm`` unlinks it).
+
+    If *d2h_stream* is provided, the D2H copy uses ``copy_()`` with
+    ``pin_memory=True`` on that stream instead of the synchronous ``.cpu()``
+    path.  The caller must synchronize *d2h_stream* after all tensors are
+    packed.
     """
     from multiprocessing import shared_memory
 
     import numpy as np
 
-    tensor = tensor.detach().cpu().contiguous()
     original_dtype = tensor.dtype
-    # NumPy does not support bfloat16; promote to float32 for the SHM
-    # transfer and record the original dtype so _tensor_from_shm can
-    # convert back.  The round-trip is lossless for bfloat16 values.
-    if original_dtype == torch.bfloat16:
-        tensor = tensor.to(torch.float32)
+    if d2h_stream is not None:
+        # Non-blocking D2H: copy on side stream to pinned CPU memory.
+        old_stream = torch.accelerator.current_stream()
+        torch.accelerator.set_stream(d2h_stream)
+        try:
+            t = tensor.detach()
+            if original_dtype == torch.bfloat16:
+                t = t.to(torch.float32)
+            cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
+            cpu.copy_(t, non_blocking=True)
+        finally:
+            torch.accelerator.set_stream(old_stream)
+        d2h_stream.synchronize()
+        tensor = cpu
+    else:
+        tensor = tensor.detach().cpu().contiguous()
+        if original_dtype == torch.bfloat16:
+            tensor = tensor.to(torch.float32)
     arr = tensor.numpy()
     nbytes = arr.nbytes
     shm = shared_memory.SharedMemory(create=True, size=nbytes)
@@ -79,14 +99,20 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     return tensor
 
 
-def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict:
+def _pack_tensor_if_large(
+    val: torch.Tensor,
+    d2h_stream: torch.Stream | None = None,
+) -> torch.Tensor | dict:
     """Replace a tensor with an SHM handle if it exceeds the threshold."""
     if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD:
-        return _tensor_to_shm(val)
+        return _tensor_to_shm(val, d2h_stream=d2h_stream)
     return val
 
 
-def _pack_value_if_large(val: object) -> object:
+def _pack_value_if_large(
+    val: object,
+    d2h_stream: torch.Stream | None = None,
+) -> object:
     """Recursively replace large tensors with SHM handles.
 
     Walks the container shapes pipelines return as ``DiffusionOutput.output``:
@@ -96,13 +122,13 @@ def _pack_value_if_large(val: object) -> object:
     the two in sync.
     """
     if isinstance(val, torch.Tensor):
-        return _pack_tensor_if_large(val)
+        return _pack_tensor_if_large(val, d2h_stream=d2h_stream)
     if isinstance(val, dict):
-        return {key: _pack_value_if_large(value) for key, value in val.items()}
+        return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()}
     if isinstance(val, list):
-        return [_pack_value_if_large(item) for item in val]
+        return [_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val]
     if isinstance(val, tuple):
-        return tuple(_pack_value_if_large(item) for item in val)
+        return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val)
     return val
 
 
@@ -119,15 +145,18 @@ def _unpack_if_shm_handle(val: object) -> object:
     return val
 
 
-def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
+def _pack_diffusion_fields(
+    output: DiffusionOutput,
+    d2h_stream: torch.Stream | None = None,
+) -> DiffusionOutput:
     if output.output is not None:
-        output.output = _pack_value_if_large(output.output)
+        output.output = _pack_value_if_large(output.output, d2h_stream=d2h_stream)
     if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
-        output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents)
+        output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream)
     if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor):
-        output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps)
+        output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream)
     if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
-        output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs)
+        output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream)
     return output
 
 
@@ -135,31 +164,37 @@ def _is_rpc_result_envelope(output: object) -> bool:
     return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
 
 
-def pack_diffusion_output_shm(output: object) -> object:
+def pack_diffusion_output_shm(
+    output: object,
+    d2h_stream: torch.Stream | None = None,
+) -> object:
     """Replace large tensors in diffusion worker outputs with SHM handles.
 
     Supports a bare ``DiffusionOutput``, a wrapper object carrying one in
     ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying
     the diffusion output in ``["result"]``, or a batch wrapper carrying
     ``RunnerOutput`` objects in ``.runner_outputs``.
+
+    If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
+    the default stream).  The caller must synchronize *d2h_stream* afterward.
     """
     if isinstance(output, DiffusionOutput):
-        return _pack_diffusion_fields(output)
+        return _pack_diffusion_fields(output, d2h_stream=d2h_stream)
 
     if _is_rpc_result_envelope(output):
         result = output.get("result")
         if isinstance(result, DiffusionOutput):
-            output["result"] = _pack_diffusion_fields(result)
+            output["result"] = _pack_diffusion_fields(result, d2h_stream=d2h_stream)
         return output
 
     result = getattr(output, "result", None)
     if isinstance(result, DiffusionOutput):
-        output.result = _pack_diffusion_fields(result)
+        output.result = _pack_diffusion_fields(result, d2h_stream=d2h_stream)
 
     runner_outputs = getattr(output, "runner_outputs", None)
     if isinstance(runner_outputs, list):
         for runner_output in runner_outputs:
-            pack_diffusion_output_shm(runner_output)
+            pack_diffusion_output_shm(runner_output, d2h_stream=d2h_stream)
     return output
 
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -3233,8 +3233,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1"
             image = image.squeeze(2)
 
-        do_denormalize = [True] * image.shape[0]
-        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+        # Denormalize + PIL moved to engine post_process_func for overlap with next request.
 
         if not return_dict:
             return (image,)

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -71,6 +71,8 @@ _STEP_COT_TEXT_LIST = "hunyuan_cot_text_list"
 _STEP_AR_KV = "hunyuan_ar_kv"
 _STEP_PROMPT_KV = "hunyuan_prompt_kv"
 
+_HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil"
+
 
 def default(val, d):
     return val if val is not None else d
@@ -333,6 +335,41 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
     return pre_process_func
 
 
+def get_hunyuan_image3_post_process_func(od_config: OmniDiffusionConfig):
+    """GPU tensor → PIL, runs outside pipeline_forward to overlap with next request."""
+    from diffusers.image_processor import VaeImageProcessor
+
+    image_processor = VaeImageProcessor()
+
+    def post_process_func(images: torch.Tensor | dict):
+        # Handle dict envelope format (upstream): {"payload": {"image": ...}, "metadata": ...}
+        if isinstance(images, dict) and "payload" in images:
+            tensor = images["payload"].get("image")
+            if tensor is None:
+                return images
+            if tensor.dim() == 3:
+                tensor = tensor.unsqueeze(0)
+            do_denormalize = [True] * tensor.shape[0]
+            images["payload"]["image"] = image_processor.postprocess(
+                tensor,
+                output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE,
+                do_denormalize=do_denormalize,
+            )
+            return images
+
+        # Legacy raw tensor format
+        if images.dim() == 3:
+            images = images.unsqueeze(0)
+        do_denormalize = [True] * images.shape[0]
+        return image_processor.postprocess(
+            images,
+            output_type=_HUNYUAN_DEFAULT_OUTPUT_TYPE,
+            do_denormalize=do_denormalize,
+        )
+
+    return post_process_func
+
+
 class HunyuanImage3Pipeline(
     HunyuanImage3PreTrainedModel,
     GenerationMixin,
@@ -2224,11 +2261,7 @@ class HunyuanImage3Pipeline(
         if hasattr(self.vae, "ffactor_temporal"):
             assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1"
             image = image.squeeze(2)
-        image = self.pipeline.image_processor.postprocess(
-            image,
-            output_type=output_type,
-            do_denormalize=[True] * image.shape[0],
-        )
+        # Postprocess deferred to engine post_process_func for overlap with next request.
 
         cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or []
         metadata = {}

vllm_omni/diffusion/registry.py

@@ -530,6 +530,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Flux2Pipeline": "get_flux2_post_process_func",
     "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
+    "HunyuanImage3Pipeline": "get_hunyuan_image3_post_process_func",
     "LingBotVideoPipeline": "get_lingbot_video_post_process_func",
     "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
@@ -540,6 +541,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "HiDreamImagePipeline": "get_hidream_image_post_process_func",
     "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func",
     "Krea2Pipeline": "get_krea2_post_process_func",
+    "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func",
 }
 
 _DIFFUSION_IR_OP_PRIORITY_FUNCS = {

vllm_omni/diffusion/sched/request_scheduler.py

@@ -50,8 +50,14 @@ class RequestScheduler(BaseScheduler):
             req_output = output.get_request_output(request_id)
             result = req_output.result if req_output is not None else None
             if result is None:
-                terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[request_id] = "No output result"
+                # Async mode: result=None with async_output_id means compute done,
+                # final output will arrive later via wait_output_ready.
+                if req_output is not None and req_output.async_output_id is not None:
+                    terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                    terminal_errors[request_id] = None
+                else:
+                    terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
+                    terminal_errors[request_id] = "No output result"
             elif result.aborted:
                 terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ABORTED
                 terminal_errors[request_id] = None

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -11,8 +11,11 @@ to DiffusionModelRunner.
 import gc
 import multiprocessing as mp
 import os
+import queue
 import signal
+import threading
 import traceback
+import uuid
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -30,6 +33,8 @@ from vllm.utils.mem_utils import GiB_bytes
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.diffusion.data import (
+    AsyncDiffusionOutput,
+    AsyncOutputKind,
     DiffusionOutput,
     OmniACK,
     OmniDiffusionConfig,
@@ -794,6 +799,21 @@ class WorkerProc:
         self.worker = self._create_worker(gpu_id, od_config, worker_extension_cls, custom_pipeline_args)
         self._running = True
 
+        self._async_output_queue: queue.Queue | None = None
+        self._async_output_thread: threading.Thread | None = None
+        if not self.od_config.step_execution:
+            self._async_output_queue = queue.Queue()
+            self._async_output_thread = threading.Thread(
+                target=self._async_output_loop,
+                daemon=True,
+                name="DiffusionAsyncOutput",
+            )
+            self._async_output_thread.start()
+
+    @staticmethod
+    def _generate_async_output_id() -> str:
+        return uuid.uuid4().hex
+
     def _create_worker(
         self,
         gpu_id: int,
@@ -813,18 +833,72 @@ class WorkerProc:
         )
         return wrapper
 
-    def _return_result(self, output: Any) -> None:
+    def _return_result(self, output: Any, rpc_id: str | None = None) -> None:
         """Reply to client, only on rank 0."""
-        if self.result_mq is not None:
-            if isinstance(output, OmniACK):
-                self.result_mq.enqueue(output)
-                return
-            try:
-                pack_diffusion_output_shm(output)
-            except Exception as e:
-                if hasattr(output, "output"):
-                    logger.warning("SHM pack failed for model output: %s", e)
+        if self.result_mq is None:
+            return
+        if isinstance(output, OmniACK):
             self.result_mq.enqueue(output)
+            return
+
+        # Async path: enqueue compute_done immediately, bg thread does D2H+SHM.
+        if not self.od_config.step_execution and isinstance(output, (DiffusionOutput, BatchRunnerOutput)):
+            async_output_id = WorkerProc._generate_async_output_id()
+            gpu_event = current_omni_platform.record_device_event()
+            self._async_output_queue.put((output, async_output_id, gpu_event))
+            msg = AsyncDiffusionOutput(
+                kind=AsyncOutputKind.COMPUTE_DONE,
+                rpc_id=rpc_id,
+                async_output_id=async_output_id,
+            )
+            self.result_mq.enqueue(msg)
+            return
+
+        # Sync path (original, or async fallback).
+        try:
+            pack_diffusion_output_shm(output)
+        except Exception as e:
+            if hasattr(output, "output"):
+                logger.warning("SHM pack failed for model output: %s", e)
+        self.result_mq.enqueue(output)
+
+    def _async_output_loop(self):
+        """Background thread: D2H + SHM packing for async diffusion output.
+
+        Uses a side stream so the D2H transfer does not block the default
+        stream where the next forward runs.
+        """
+        device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id)
+        d2h_stream = torch.Stream(device=device)
+        while self._running:
+            output, async_output_id, gpu_event = self._async_output_queue.get()
+            try:
+                # Cross-stream ordering: wait for default stream to finish
+                # writing the output tensors before the side stream reads.
+                if gpu_event is not None:
+                    d2h_stream.wait_event(gpu_event)
+                pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
+                d2h_stream.synchronize()
+
+                self.result_mq.enqueue(
+                    AsyncDiffusionOutput(
+                        kind=AsyncOutputKind.OUTPUT_READY,
+                        async_output_id=async_output_id,
+                        output=output,
+                    )
+                )
+            except Exception:
+                logger.exception(
+                    "Async output packing failed for id '%s'; sending error",
+                    async_output_id,
+                )
+                self.result_mq.enqueue(
+                    AsyncDiffusionOutput(
+                        kind=AsyncOutputKind.OUTPUT_READY,
+                        async_output_id=async_output_id,
+                        error="Background D2H/SHM packing failed",
+                    )
+                )
 
     def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
         if not torch.distributed.is_initialized():
@@ -936,13 +1010,25 @@ class WorkerProc:
             # Route message based on type
             elif isinstance(msg, dict) and msg.get("type") == "rpc":
                 try:
+                    rpc_id = msg.get("rpc_id")
                     result, should_reply = self._execute_rpc(msg)
                     if should_reply:
-                        self._return_result(result)
+                        self._return_result(result, rpc_id=rpc_id)
                 except Exception as e:
                     logger.error(f"Error processing RPC: {e}", exc_info=True)
                     if self.result_mq is not None:
-                        self._return_result({"status": "error", "error": str(e)})
+                        if rpc_id is not None:
+                            # Async RPC: must complete the executor''s pending
+                            # future so collective_rpc() doesn''t hang.
+                            self.result_mq.enqueue(
+                                AsyncDiffusionOutput(
+                                    kind=AsyncOutputKind.RPC_RESULT,
+                                    rpc_id=rpc_id,
+                                    error=str(e),
+                                )
+                            )
+                        else:
+                            self._return_result({"status": "error", "error": str(e)})
 
             elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                 logger.info("Worker %s: Received shutdown message", self.gpu_id)

vllm_omni/diffusion/worker/utils.py

@@ -183,6 +183,7 @@ class RunnerOutput(BaseRunnerOutput):
     step_index: int | None = None
     finished: bool = False
     result: DiffusionOutput | None = None
+    async_output_id: str | None = None
 
     def get_request_output(self, request_id: str) -> RunnerOutput | None:
         return self if self.request_id == request_id else None

vllm_omni/platforms/cuda/platform.py

@@ -273,6 +273,17 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
     def synchronize(cls) -> None:
         torch.accelerator.synchronize()
 
+    @classmethod
+    def record_device_event(cls) -> torch.Event | None:
+        """Record a device event on the default stream to mark tensor readiness."""
+        try:
+            event = torch.Event()
+            event.record()
+            return event
+        except Exception:
+            logger.warning("Failed to record device event for cross-stream sync")
+            return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         free, _ = torch.cuda.mem_get_info(device)

vllm_omni/platforms/interface.py

@@ -190,6 +190,24 @@ class OmniPlatform(Platform):
     def synchronize(cls) -> None:
         raise NotImplementedError
 
+    # ── Async diffusion output: cross-stream sync ──
+
+    @classmethod
+    def record_device_event(cls):
+        """Record a device event on the default stream to mark tensor readiness.
+
+        On platforms where distributed communication (e.g. HCCL) may use
+        internal streams not visible to the default stream, this method
+        should synchronize the default stream before recording the event
+        to ensure the event captures all completed work including
+        cross-device communication results.
+
+        Returns ``None`` by default so that platforms without a native
+        implementation (ROCm, XPU, MUSA) fall through to a safe no-op.
+        Override in platform subclasses to provide real event support.
+        """
+        return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError

vllm_omni/platforms/musa/platform.py

@@ -152,6 +152,16 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         """Synchronize all MUSA operations."""
         torch.musa.synchronize()
 
+    @classmethod
+    def record_device_event(cls) -> torch.Event | None:
+        try:
+            event = torch.musa.Event()
+            event.record()
+            return event
+        except Exception:
+            logger.warning("Failed to record MUSA device event for cross-stream sync")
+            return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         free, _ = torch.musa.mem_get_info(device)

vllm_omni/platforms/npu/platform.py

@@ -175,6 +175,25 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     def synchronize(cls) -> None:
         torch.npu.synchronize()
 
+    @classmethod
+    def record_device_event(cls) -> torch.Event | None:
+        """Record a NPU event on the default stream to mark tensor readiness.
+
+        On NPU/Ascend with HCCL, distributed communication may use internal
+        streams not visible to the default stream. Synchronize the default
+        stream first so that HCCL results are written back before we record
+        the event, ensuring d2h_stream.wait_event() captures the complete
+        output data.
+        """
+        try:
+            torch.npu.current_stream().synchronize()
+            event = torch.npu.Event()
+            event.record()
+            return event
+        except Exception:
+            logger.warning("Failed to record NPU event for cross-stream sync")
+            return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         free, _ = torch.npu.mem_get_info(device)

vllm_omni/platforms/rocm/platform.py

@@ -145,6 +145,16 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
     def synchronize(cls) -> None:
         torch.accelerator.synchronize()
 
+    @classmethod
+    def record_device_event(cls) -> torch.Event | None:
+        try:
+            event = torch.Event()
+            event.record()
+            return event
+        except Exception:
+            logger.warning("Failed to record device event for cross-stream sync")
+            return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         free, _ = torch.cuda.mem_get_info(device)

vllm_omni/platforms/xpu/platform.py

@@ -99,6 +99,16 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
     def synchronize(cls) -> None:
         torch.xpu.synchronize()
 
+    @classmethod
+    def record_device_event(cls) -> torch.Event | None:
+        try:
+            event = torch.xpu.Event()
+            event.record()
+            return event
+        except Exception:
+            logger.warning("Failed to record XPU device event for cross-stream sync")
+            return None
+
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         free, _ = torch.xpu.mem_get_info(device)