vllm-omni/docs/design/feature/async_diffusion_output.md

Async Diffusion Output

Table of Contents

1. Overview
2. Performance
3. Architecture
4. Model Coverage
5. Configuration
6. 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_ids (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