From dba05c3517e77b1d3ec85a484ef5286834f8a600 Mon Sep 17 00:00:00 2001 From: zzh <67720593+BLANKETusers@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:21:55 +0800 Subject: [PATCH] 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 | 174 ++++++++++++ tests/diffusion/test_async_output_worker.py | 193 +++++++++++++ tests/diffusion/test_ipc_async.py | 159 +++++++++++ .../test_multiproc_engine_concurrency.py | 13 +- tests/diffusion/test_result_pump.py | 256 +++++++++++++++++ tests/platforms/test_record_device_event.py | 79 ++++++ vllm_omni/diffusion/data.py | 32 +++ vllm_omni/diffusion/diffusion_engine.py | 15 +- .../diffusion/executor/multiproc_executor.py | 261 ++++++++++++++++-- vllm_omni/diffusion/ipc.py | 83 ++++-- .../hunyuan_image3_transformer.py | 3 +- .../hunyuan_image3/pipeline_hunyuan_image3.py | 43 ++- vllm_omni/diffusion/registry.py | 2 + .../diffusion/sched/request_scheduler.py | 10 +- .../diffusion/worker/diffusion_worker.py | 110 +++++++- vllm_omni/diffusion/worker/utils.py | 1 + vllm_omni/platforms/cuda/platform.py | 11 + vllm_omni/platforms/interface.py | 18 ++ vllm_omni/platforms/musa/platform.py | 10 + vllm_omni/platforms/npu/platform.py | 19 ++ vllm_omni/platforms/rocm/platform.py | 10 + vllm_omni/platforms/xpu/platform.py | 10 + 22 files changed, 1440 insertions(+), 72 deletions(-) create mode 100644 docs/design/feature/async_diffusion_output.md create mode 100644 tests/diffusion/test_async_output_worker.py create mode 100644 tests/diffusion/test_ipc_async.py create mode 100644 tests/diffusion/test_result_pump.py create mode 100644 tests/platforms/test_record_device_event.py diff --git a/docs/design/feature/async_diffusion_output.md b/docs/design/feature/async_diffusion_output.md new file mode 100644 index 000000000..a81ed6811 --- /dev/null +++ b/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 diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py new file mode 100644 index 000000000..8cda3be1c --- /dev/null +++ b/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" diff --git a/tests/diffusion/test_ipc_async.py b/tests/diffusion/test_ipc_async.py new file mode 100644 index 000000000..c54d2f793 --- /dev/null +++ b/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, + ) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index f78f8d44b..10fe7e86c 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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] diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py new file mode 100644 index 000000000..fe49df47e --- /dev/null +++ b/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 diff --git a/tests/platforms/test_record_device_event.py b/tests/platforms/test_record_device_event.py new file mode 100644 index 000000000..a979ca940 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d3880f900..ca9d48d96 100644 --- a/vllm_omni/diffusion/data.py +++ b/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.""" diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 4d37843d7..44ca9da89 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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) diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index c7ca34153..84a460c53 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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 = [] diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index cd83e718c..ca6bc858d 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/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 diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index f749d73d3..877fd1c3f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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,) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 095a420c8..36e69302e 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/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 = {} diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index aaedc7b00..1fbddcb96 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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 = { diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index 9e7442d3a..ec6b2aa8e 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b7450038a..58c8fd1eb 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index b49388fee..40ed550d0 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/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 diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 5eb7406a5..642d8ec02 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 9969849d9..99dafa2f2 100644 --- a/vllm_omni/platforms/interface.py +++ b/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 diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 818338da4..c71b5a70b 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/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) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 240f9d576..bb0e1b436 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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) diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index f44bb86fa..cf72baccb 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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) diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 26d3dccde..aa26cd82c 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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)