Motivation.

This feature offloads the D2H (Device-to-Host) copy and SHM packing—performed after the Worker GPU forward pass—from the Worker's main thread to a background daemon thread. This allows the GPU's default stream to immediately start the forward pass for the next request, thereby eliminating GPU bubbles.

Proposed Change.

Async Diffusion Output

1. Overview

After the Worker completes GPU forward, it must synchronously execute D2H (Device→Host) copy and SHM packing on the main thread, during which the GPU default stream sits idle, creating a GPU bubble. In multi-request concurrent scenarios, the D2H/packing of the previous request blocks the forward launch of the next request, reducing overall throughput.

This feature moves the D2H copy and SHM packing from the Worker main thread to a background daemon thread (side CUDA stream), so the GPU default stream can immediately start the next request's forward, achieving overlap between the previous request's output D2H/packing and the next request's forward execution, eliminating the GPU bubble.

2. Scope & Objectives

Goals

• Eliminate GPU bubble in request-mode: After Worker forward completes, immediately return compute_done; the GPU default stream can start the next forward, while D2H/packing executes asynchronously on the side stream.
• Automatically enabled, zero configuration: When step_execution=False (default, request-mode), the async output path is automatically enabled; when step_execution=True (step-mode), the original synchronous path is used. No extra configuration switch needed.
• Performance improvement: On HunyuanImage-3.0 (TP4), 768×768 resolution QPS improved by +1.95%, 1024×1024 resolution QPS improved by +0.60%.
• Compatibility: Transparent to request-mode models — existing models benefit without modification; step-mode behavior is completely unchanged.

Feature Goals

• Step-mode async output adaptation: Currently not supported for step_execution=True models. Step-mode requires synchronous intermediate results per step; adaptation is more complex and deferred to a future phase.
• Intermediate step async D2H: Even when adapting step-mode in the future, only the final step's async D2H is planned; intermediate steps will remain synchronous.

3. Design

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       │
                    └─────────────────────────────────────────────────────┘

Sequence Diagram (request-mode, async path)

sequenceDiagram
    autonumber
    participant Client as Online Client
    participant StageA as Stage thread A
    participant StageB as Stage thread B
    participant Engine as DiffusionEngine
    participant Executor as MultiprocDiffusionExecutor
    participant Pump as result_pump
    participant WorkerLoop as WorkerProc.worker_busy_loop
    participant Worker as DiffusionWorker
    participant PackThread as async output thread
    participant MQ as result_mq

    Client->>StageA: request 1
    StageA->>Engine: step(req1)
    Engine->>Engine: acquire _rpc_lock
    Engine->>Executor: execute_request(req1)
    Executor->>Executor: _next_rpc_id() → rpc_id
    Executor->>WorkerLoop: broadcast rpc execute_model(req1, rpc_id)
    WorkerLoop->>Worker: execute_model(req1)
    Worker-->>WorkerLoop: DiffusionOutput(device tensor)
    WorkerLoop->>PackThread: submit_async_output(async_output_id, output)
    WorkerLoop->>MQ: enqueue(AsyncDiffusionOutput(kind=COMPUTE_DONE, rpc_id, async_output_id))
    WorkerLoop-->>WorkerLoop: dequeue next request
    MQ-->>Pump: dequeue()
    Pump-->>Executor: resolve rpc_futures[rpc_id]
    Executor-->>Engine: RunnerOutput(result=None, async_output_id=...)
    Engine->>Engine: scheduler.update_from_output()<br/>release running slot
    Engine->>Engine: release _rpc_lock

    par req1 output packing
        PackThread->>PackThread: wait gpu_event
        PackThread->>PackThread: pack_diffusion_output_shm()<br/>D2H + SHM write on side stream
        PackThread->>MQ: enqueue(AsyncDiffusionOutput(kind=OUTPUT_READY, async_output_id))
        MQ-->>Pump: dequeue()
        Pump-->>Executor: resolve output_futures[async_output_id]
        StageA->>Executor: wait_output_ready(async_output_id)
        Executor-->>StageA: DiffusionOutput
        StageA->>StageA: post_process_func(output)
    and req2 forward can start
        Client->>StageB: request 2
        StageB->>Engine: step(req2)
        Engine->>Engine: acquire _rpc_lock
        Engine->>Executor: execute_request(req2)
        Executor->>WorkerLoop: broadcast rpc execute_model(req2)
        WorkerLoop->>Worker: execute_model(req2)
    end

    StageA-->>Client: final result for req1

    Note over WorkerLoop,PackThread: D2H/SHM packing leaves worker main loop.<br/>req2 forward can overlap with req1 output packing.

Optional API & Interface Changes

New Classes

Class	File	Description
AsyncOutputKind	vllm_omni/diffusion/data.py	Enum: RPC_RESULT, COMPUTE_DONE, OUTPUT_READY
AsyncDiffusionOutput	vllm_omni/diffusion/data.py	Protocol envelope for result_mq; kind field routes messages

class AsyncOutputKind(Enum):
    RPC_RESULT = "rpc_result"       # Ordinary RPC return (including error propagation)
    COMPUTE_DONE = "compute_done"   # Forward finished, GPU can start next request
    OUTPUT_READY = "output_ready"   # D2H/SHM packing finished, final output available

@dataclass
class AsyncDiffusionOutput:
    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

Modified Classes

Class	File	Change
DiffusionOutput	vllm_omni/diffusion/data.py	New field async_output_id: str | None = None — async output claim token
RunnerOutput	vllm_omni/diffusion/data.py	New field async_output_id: str | None = None — marks "compute done, output pending"

New Methods

Method	File	Description
_generate_async_output_id()	vllm_omni/diffusion/worker/diffusion_worker.py	Static method; generates UUID claim token
_async_output_loop()	vllm_omni/diffusion/worker/diffusion_worker.py	Daemon thread main loop: dequeue from _async_output_queue → wait_event → side-stream D2H+SHM packing → enqueue output_ready
record_device_event()	vllm_omni/platforms/interface.py	Base class returns None (safe no-op); subclasses override for CUDA/NPU/ROCm/XPU/MUSA
get_hunyuan_image3_post_process_func()	vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py	Postprocess function moved from Worker to Engine side

New Attributes

Attribute	File	Description
_async_output_queue	diffusion_worker.py	queue.Queue; main thread puts (output, async_output_id, gpu_event), bg thread consumes
_async_output_thread	diffusion_worker.py	Daemon thread object
_rpc_id_counter + _rpc_id_lock	multiproc_executor.py	RPC ID counter + lock, Path 1 only
_rpc_futures	multiproc_executor.py	dict[str, Future[AsyncDiffusionOutput]]; wait compute_done by rpc_id
_output_futures	multiproc_executor.py	dict[str, Future[DiffusionOutput]]; wait final output by async_output_id
_completed_outputs	multiproc_executor.py	dict[str, Future[DiffusionOutput]]; cache for output_ready arriving before wait
_batch_split_map	multiproc_executor.py	dict[str, dict[str, str]]; batch_id → {per_req_id: request_id}
_futures_lock	multiproc_executor.py	RLock; protects above four dicts
_pump_stop	multiproc_executor.py	threading.Event; shutdown signal for pump thread
_sync_result_buffer	multiproc_executor.py	queue.Queue; pump puts non-async messages here for Path 2

Modified Methods

Method	File	Change
_return_result()	diffusion_worker.py	New async branch: when step_execution=False and output is DiffusionOutput/BatchRunnerOutput, generate async_output_id + gpu_event, put to _async_output_queue, enqueue compute_done, return immediately
worker_busy_loop()	diffusion_worker.py	Extract rpc_id and pass to _return_result(); add async RPC error propagation via AsyncDiffusionOutput(kind=RPC_RESULT, error=...)
_dequeue_one_with_failure_polling()	multiproc_executor.py	New dispatch: request-mode reads from _sync_result_buffer; step-mode reads from _result_mq (original logic)
execute_request()	multiproc_executor.py	New compute_done branch: construct RunnerOutput(finished=True, result=None, async_output_id=...)
execute_batch()	multiproc_executor.py	New compute_done branch: split batch-level async_output_id into per-request via _batch_split_map
_result_pump()	multiproc_executor.py	Handle OUTPUT_READY with batch split logic via _batch_split_map
wait_output_ready()	multiproc_executor.py	Return cached completed Future directly
shutdown()	multiproc_executor.py	Add _batch_split_map.clear()
step()	diffusion_engine.py	Check async_output_id → await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT)
step_streaming()	diffusion_engine.py	Same async wait inside async for loop
add_req_and_wait_for_response()	diffusion_engine.py	fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT) blocking wait
_finalize_finished_request()	diffusion_engine.py	New branch: result=None but async_output_id present → return DiffusionOutput(async_output_id=...) placeholder
update_from_output()	request_scheduler.py	result=None + async_output_id present → FINISHED_COMPLETED (not FINISHED_ERROR)
_tensor_to_shm()	ipc.py	New d2h_stream param: side-stream path uses pin_memory + copy_(non_blocking=True); original .cpu() when absent
_pack_tensor_if_large()	ipc.py	Pass through d2h_stream
_pack_value_if_large()	ipc.py	Pass through d2h_stream
_pack_diffusion_fields()	ipc.py	Pass through d2h_stream
pack_diffusion_output_shm()	ipc.py	Pass through d2h_stream (entry function)

Modified Pipeline Files (postprocess removal)

File	Change
pipeline_hunyuan_image3.py	HunyuanImage3Pipeline.__call__: postprocess commented out, output=image (return raw GPU tensor)
hunyuan_image3_transformer.py	HunyuanImage3Text2ImagePipeline.__call__: postprocess commented out
registry.py	_DIFFUSION_POST_PROCESS_FUNCS registers HunyuanImage3Pipeline and HunyuanImage3ForCausalMM → get_hunyuan_image3_post_process_func

Key Technical Decisions

Decision 1: Two-phase message protocol (COMPUTE_DONE + OUTPUT_READY)

• Decision: Split one request's completion into compute_done and output_ready two phases via AsyncDiffusionOutput.kind, using the same physical result_mq.
• Alternatives: (A) Use two separate queues for compute completion and output delivery; (B) Use a single message type with a flag field.
• Rationale: Single queue preserves the existing result_mq infrastructure and ordering guarantees. The kind field provides clean dispatch without requiring the pump thread to manage two consumer loops. Two separate queues would complicate ordering and error handling.

Decision 2: Pump thread as sole reader of result_mq

• Decision: In request-mode, the pump thread is the sole reader of result_mq. collective_rpc() cannot directly dequeue; non-async messages are buffered in _sync_result_buffer.
• Alternatives: (A) Multiple readers with locks on result_mq; (B) Let collective_rpc() directly read result_mq when expecting async messages.
• Rationale: Single reader eliminates race conditions on result_mq and ensures message ordering. The _sync_result_buffer indirection is minimal overhead for Path 2 RPCs while keeping the pump's dispatch logic simple.

Decision 3: Two-path dispatch in collective_rpc()

• Decision: Path 1 (execute_model/execute_model_batch) generates rpc_id and waits for pump-delivered compute_done via Future; Path 2 (other RPCs) reads from _sync_result_buffer or _result_mq.
• Alternatives: (A) All RPCs go through the same async path; (B) Per-method configuration.
• Rationale: Only execute_model/execute_model_batch benefit from the compute_done/output_ready split (they produce GPU tensors requiring D2H). Other RPCs (sleep/wake/profile) return small results with no D2H overhead; forcing them through the async path adds unnecessary latency and complexity. The whitelist approach is simple and explicit.

Decision 4: Postprocess moved from Worker to Engine

• Decision: Move denormalize→PIL conversion from Worker pipeline to Engine side via registered post_process_func.
• Alternatives: (A) Keep postprocess in Worker bg thread after D2H; (B) Skip postprocess entirely.
• Rationale: Postprocess is CPU-intensive. If kept in Worker bg thread, it increases packing time and compresses the overlap benefit. Moving it to Engine side reduces bg thread work to pure D2H+SHM packing, maximizing overlap with the next forward. The postprocess runs on Engine side after wait_output_ready() returns, which is already outside the critical GPU path.

Decision 5: record_device_event() base class returns None

• Decision: OmniPlatform.record_device_event() base class returns None (safe no-op) instead of raising NotImplementedError.
• Alternatives: (A) Raise NotImplementedError to force subclasses to implement; (B) Use a separate capability flag.
• Rationale: Safe degradation is preferred over hard failures. If a platform doesn't implement record_device_event(), the bg thread will skip the event wait and proceed with D2H immediately — this may cause data races but won't crash. This allows incremental platform support without breaking existing deployments.

Dependencies & Risks

• Memory pressure from pending outputs: Background D2H holds output tensor references. If pending outputs accumulate (e.g., Engine side is slow to consume), GPU memory and host memory pressure will increase. Mitigation: _ASYNC_OUTPUT_TIMEOUT provides a hard timeout; if exceeded, the request fails rather than blocking indefinitely.
• Batch split correctness: execute_batch path uses _batch_split_map to split batch OUTPUT_READY into per-request results. The split logic depends on DiffusionOutput.get_request_output() method. If the batch output structure changes, the split logic must be updated accordingly.
• RPC error propagation: Worker RPC exceptions must be propagated via AsyncDiffusionOutput(kind=RPC_RESULT, error=...) to avoid collective_rpc() pending Future hanging forever. This is handled in worker_busy_loop() exception handling.
• _completed_outputs stores Future, not bare DiffusionOutput: wait_output_ready() returns the cached completed Future directly. Callers must not assume they receive a DiffusionOutput — they must call .result() on the Future.
• Pump thread lifecycle: Pump thread is started in request-mode and stopped via _pump_stop event during shutdown(). If pump crashes, all collective_rpc() calls will hang. Mitigation: _dequeue_one_with_failure_polling() checks is_failed and _closed on timeout.
• Step-mode compatibility: Step-mode is completely unaffected — pump and Worker bg thread are not started. No risk of regression for step-mode models.

Future Work: Step-Mode Async Output Adaptation

Does Step-Mode Need Async Output?

• Intermediate steps return result=None (no D2H): no GPU bubble. Only the final step's post_decode() + D2H/packing causes a bubble.
• Overlap opportunity: in multi-request scenarios, the final step's D2H/packing of request N can overlap with the first denoise step of request N+1. Single-request scenarios see no benefit.
• Conclusion: step-mode benefits from async output only in multi-request concurrency, limited to the final step's D2H overlap with the next request's forward.

Key Changes Needed

1. execute_stepwise → Path 1: Add "execute_stepwise" to the collective_rpc() Path 1 whitelist (or replace whitelist with step_execution-based check).
2. Start pump thread in step-mode: Enable _result_pump regardless of step_execution so COMPUTE_DONE/OUTPUT_READY dispatch works.
3. _return_result() for step-mode: Only the final step (result=DiffusionOutput) needs async path; intermediate steps (result=None) remain synchronous.
4. _busy_loop async wait: When runner_output contains async_output_id, let _busy_loop wait for output_ready before _emit_outputs(), while update_from_output() advances immediately.
5. post_decode() externalization: Move post_decode() from Worker to Engine (same pattern as request-mode postprocess externalization), so bg thread only handles D2H+SHM packing.

Open Questions （Verification needed）

• Streaming output (streaming_output=True): Each chunk_denoise_completed step also needs D2H. The D2H of step N can overlap with step N+1's forward (same pattern as request-mode), but each chunk output is smaller than the full output, so the absolute time saved is less.

Feedback Period.

Currently, the benefits of this feature have been verified on the HunyuanImage3 and QwenImage models. In theory, any model configured with step_execution=False (i.e., in request-mode) should see similar improvements. We'd appreciate it if everyone could help verify the performance gains on other models.

CC List.

No response

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.