vllm-omni/docs/design/module/diffusion/diffusion_runtime.md

title, kind, status, owners, primary_code_paths, related_code_paths, depends_on, validation_paths, upstream_refs, last_reviewed, last_verified_commit

title	kind	status	owners	primary_code_paths	related_code_paths	depends_on	validation_paths	upstream_refs	last_reviewed	last_verified_commit
Diffusion Runtime	module	draft	@Isotr0py @princepride @SamitHuang @fhfuih	vllm_omni/diffusion/diffusion_engine.py vllm_omni/diffusion/request.py vllm_omni/diffusion/data.py vllm_omni/diffusion/sched/** vllm_omni/diffusion/executor/** vllm_omni/diffusion/worker/**	vllm_omni/diffusion/stage_diffusion_client.py vllm_omni/diffusion/stage_diffusion_proc.py vllm_omni/diffusion/inline_stage_diffusion_client.py vllm_omni/diffusion/ipc.py vllm_omni/diffusion/hooks/** vllm_omni/diffusion/postprocess/**	../engine_orchestration.md ../stage_runtime.md ../input_output_modality_contracts.md diffusion_model_integration.md	tests/diffusion/test_diffusion_engine.py tests/diffusion/test_diffusion_engine_cleanup.py tests/diffusion/test_diffusion_engine_dummy_run.py tests/diffusion/test_diffusion_engine_rpc_routing.py tests/diffusion/test_diffusion_scheduler.py tests/diffusion/test_diffusion_model_runner.py tests/diffusion/test_diffusion_worker.py tests/diffusion/test_multiproc_engine_concurrency.py tests/diffusion/test_uniproc_executor.py tests/diffusion/test_result_pump.py tests/diffusion/test_async_output_worker.py tests/diffusion/test_diffusion_ipc.py tests/diffusion/test_diffusion_streaming_output.py tests/diffusion/test_stage_diffusion_proc.py tests/diffusion/test_inline_stage_diffusion_client.py	diffusers.DiffusionPipeline	2026-09-02	e51fe6ec1b

Diffusion runtime

The diffusion runtime owns request admission, scheduling, execution, progress, output, cancellation, and cleanup inside a diffusion stage.

It is easiest to think of the runtime as a small control loop:

1. the scheduler chooses ready requests;
2. the executor sends that choice to the worker layer;
3. each worker manages its device and delegates model work to a runner;
4. the runner calls the model pipeline and returns per-request results; and
5. the engine feeds those results back to the scheduler and output streams.

!!! note

This page describes the runtime plumbing. Model loading and pipeline
contracts belong to [Diffusion model integration](diffusion_model_integration.md).
Batch compatibility and step-state details belong to
[Diffusion continuous batching](continuous_batching.md).

Goals and non-goals

The runtime aims to keep request policy separate from device execution. It also keeps request identity stable from admission through the final output, including cancellation and failures.

This module does not choose the next Omni stage, place stage replicas, or define a model's denoising algorithm:

• cross-stage routing belongs to Engine orchestration;
• process placement and replica lifecycle belong to Stage runtime; and
• model-specific work belongs to Diffusion model integration.

Runtime at a glance

The stage owns DiffusionEngine and a DiffusionExecutor backend. The inline stage client keeps that stack in the caller process; a process-backed stage runs it inside StageDiffusionProc.

The executor then chooses how the worker runs:

• uni (default when num_gpus == 1): UniProcDiffusionExecutor builds one in-process worker. There is no worker subprocess and no shared-memory RPC.
• mp (default when num_gpus > 1, or when set explicitly): MultiprocDiffusionExecutor starts one WorkerProc per device.

Either way, each worker owns its device, distributed state, DiffusionWorker, DiffusionModelRunner, and pipeline instance.

flowchart TB
    client["Stage client"]

    subgraph stage["Diffusion stage process"]
        stageProc["StageDiffusionProc<br/>stage transport and lifecycle"]
        engine["DiffusionEngine<br/>admission and control loop"]
        scheduler["RequestScheduler or StepScheduler<br/>request state and policy"]
        executor["DiffusionExecutor<br/>uni or mp backend"]
        streams["Per-request output streams"]

        stageProc --> engine
        engine --> scheduler
        scheduler -->|"DiffusionSchedulerOutput"| engine
        engine --> executor
        engine --> streams
    end

    subgraph workers["Worker layer"]
        workerProc["WorkerProc<br/>mp IPC loop only"]
        worker["DiffusionWorker<br/>device and distributed setup"]
        runner["DiffusionModelRunner<br/>model state and execution"]
        pipeline["Diffusion pipeline<br/>model-specific computation"]

        workerProc --> worker
        worker --> runner --> pipeline
    end

    client -->|"process-backed: ZMQ"| stageProc
    client -.->|"inline: direct call"| engine
    executor -->|"mp: broadcast / result queues"| workerProc
    executor -.->|"uni: in-process call"| worker
    streams -->|"process-backed"| stageProc
    stageProc -->|"ZMQ results"| client
    streams -.->|"inline"| client

!!! info "Stage process vs worker process"

A deployed stage may already run in a `StageDiffusionProc` child process.
That is separate from the optional `mp` worker processes. With `uni`, the
worker lives in the same process as the engine.

Component responsibilities

Component	Owns	Does not own
DiffusionEngine	Admission, the busy loop, scheduler/executor coordination, RPC ordering, output streams, warmup, abort routing	Batch policy details, device setup, model computation
BaseScheduler	Request states, waiting/running sets, capacity, compatibility, terminal transitions	IPC, worker calls, output formatting
RequestScheduler	Complete-request waves and optional admission delay for request-level batching	Denoising-step progress
StepScheduler	Per-request denoising progress and step-wise completion	Pipeline tensors and model state
DiffusionExecutor	Execution backend contract, worker RPC, health, shutdown	Admission and request-state transitions
UniProcDiffusionExecutor	Single in-process worker for num_gpus == 1; no IPC or async output pump	Multi-GPU execution
MultiprocDiffusionExecutor	Worker processes, shared-memory message queues, result dispatch, worker monitoring	In-process single-GPU path
WorkerProc	One mp worker process's IPC loop and reply rules	Scheduling policy; unused by uni
DiffusionWorker	Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation	Request admission
DiffusionModelRunner	Pipeline loading, request-local model state, cache/compile setup, request or step execution	Queueing and cross-stage routing

The boundary between policy and execution is DiffusionSchedulerOutput. It contains newly admitted request payloads, IDs for requests whose runner state is already cached, finished IDs, optional KV-prefetch work, and per-request diffusion_kv_metadata on NewRequestData envelopes when paged Diffusion KV is enabled. Executors and workers consume this output; they do not decide what should run next.

BaseScheduler.initialize() sets max_num_running_reqs from od_config.max_num_seqs (default 1). The engine may override that capacity for distributed layerwise offload with AllGather DP concurrency.

One scheduling tick

Async requests share one engine busy loop. Worker calls and control RPCs pass through that loop, so they cannot race on the executor transport.

sequenceDiagram
    participant Caller
    participant Engine as DiffusionEngine
    participant Scheduler
    participant Executor
    participant Worker
    participant Runner

    Caller->>Engine: add request
    Engine->>Scheduler: add_request(request)
    Engine->>Scheduler: schedule()
    Scheduler-->>Engine: DiffusionSchedulerOutput
    Engine->>Executor: execute_batch() or execute_step()
    Executor->>Worker: run worker method
    Worker->>Runner: execute model or one step
    Runner-->>Worker: RunnerOutput(s)
    Worker-->>Executor: result
    Executor-->>Engine: BaseRunnerOutput
    Engine->>Scheduler: update_from_output(...)
    Scheduler-->>Engine: finished request IDs
    Engine-->>Caller: chunk or final output

The engine catches request execution errors and turns them into per-request error outputs. A dead worker group is different: the executor marks itself failed, health checks raise EngineDeadError, and the owning stage client handles the stage-level failure.

Execution modes

The engine resolves one mode at startup and binds the matching scheduler and executor call.

Mode	Scheduler	Executor call	Runner path
Request batch	RequestScheduler	execute_batch()	execute_model() for one request, or execute_model_batch() for a fused batch
Step batch	StepScheduler	execute_step()	execute_stepwise()

Request mode runs a complete pipeline forward for each scheduled wave. A single-request wave is the conservative path. A multi-request wave usually uses fused execute_model_batch() and requires the pipeline to declare request-batch support.

Distributed layerwise offload with AllGather (DLO DP concurrency) is a separate multiproc dispatch path. It activates when data_parallel_size > 1, enable_distributed_layerwise_offload is set, and dlo_use_allgather is true. The engine then sets dp_concurrent = True and raises scheduler.max_num_running_reqs to dp_size, overriding the max_num_seqs value from initialize(). The scheduler still emits a multi-request DiffusionSchedulerOutput; the multiproc executor routes that wave through execute_request() instead of fused execute_model_batch().

Optional admission coalescing can fill a dp_size wave: when request_batch_max_wait_ms > 0, RequestScheduler may wait briefly before the first schedule of a wave (under dp_concurrent, the stable window is min(0.3s, wait/2)). With the default request_batch_max_wait_ms == 0, there is no wait.

Before dispatch, MultiprocDiffusionExecutor.execute_request() rejects waves whose requests differ in sampling-parameter compatibility or extra_args (AllGather requires every DP rank to follow the same forward schedule). Each worker rank then picks one envelope from the wave: req[dp_rank % len(req)], so every rank enters the same collective while computing different requests.

Step mode keeps StepRequestState in the runner. New requests run prepare_encode() once; every tick runs denoise_step() and step_scheduler(); completed requests run post_decode(). The scheduler keeps only lifecycle and progress metadata—it does not hold model tensors.

Step mode is also the path for streaming diffusion output. Enabling streaming_output turns on step execution if needed. Chunk-capable pipelines can emit intermediate outputs through the same request stream; final-only step pipelines still emit only the final result. If the pipeline does not implement step execution, initialization fails.

!!! tip

For user-facing flags and examples, see
[Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md).
That guide explains how to select a mode; this page explains what happens
after the choice is made.

Request state and identity

The scheduler is the source of truth for an admitted request. Its normal state flow is:

stateDiagram-v2
    [*] --> WAITING: add_request
    WAITING --> RUNNING: schedule
    WAITING --> FINISHED_ABORTED: abort
    RUNNING --> FINISHED_ABORTED: abort
    RUNNING --> FINISHED_COMPLETED: successful output
    RUNNING --> FINISHED_ERROR: failed or missing output
    FINISHED_COMPLETED --> [*]: deliver and remove state
    FINISHED_ABORTED --> [*]: deliver and remove state
    FINISHED_ERROR --> [*]: deliver and remove state

PREEMPTED and BaseScheduler.preempt_request() exist on the scheduler API, but the engine loop does not call them today—only scheduler unit tests do.

request_id ties together scheduler state, runner state, executor results, and the output queue. Batch positions are temporary and must never replace request identity when results are mapped back.

Only compatible requests enter the same running wave. The scheduler compares a mode-specific sampling-parameter key and stops at the first incompatible waiting request. This is intentionally conservative and can cause head-of-line blocking.

Executor and IPC

DiffusionExecutor is the backend interface. Built-in backends are selected by distributed_executor_backend:

Backend	When chosen	Worker layout
uni	Default for num_gpus == 1, or set explicitly	One in-process worker
mp	Default for num_gpus > 1, or set explicitly	One worker process per device

Ray and external-launcher diffusion backends are not implemented. A custom DiffusionExecutor subclass or import path is also accepted.

Uniproc backend

UniProcDiffusionExecutor constructs WorkerWrapperBase in the engine process and calls worker methods directly. It avoids a second model load, MessageQueue rings, ZMQ IPC sockets, and /dev/shm tensor packing. RPC timeouts are accepted for interface parity but are not enforced: a hung worker blocks the calling thread. Sticky accelerator faults mark the executor dead through failure callbacks; ordinary request errors stay per-request.

Multiproc backend

MultiprocDiffusionExecutor:

1. creates a broadcast queue shared by all workers;
2. spawns one worker process per configured device;
3. waits until every worker reports ready;
4. sends generation and control calls as RPC messages;
5. applies rank-aware reply rules so only expected ranks respond; and
6. monitors worker process sentinels and fails the executor if a worker dies.

In request mode on this backend, each async output produces two correlated messages keyed by async_output_id:

• COMPUTE_DONE — the forward finished and the device can start the next request.
• OUTPUT_READY — background D2H/SHM packing finished and the final output is ready.

The worker queues background packing work before it enqueues COMPUTE_DONE, but either message may reach the executor first. The result pump and execute_batch() must handle both orderings. Step mode keeps the synchronous result path and does not start those pumps.

See Async diffusion output for the multiproc request-mode timeline.

Worker and runner boundary

On the mp path, WorkerProc receives messages and calls methods through WorkerWrapperBase. On the uni path, the executor owns that wrapper directly. In both cases, DiffusionWorker handles the parts tied to a device: distributed initialization, model-runner construction, LoRA activation, profiling, and memory sleep/wake. It delegates actual model work to DiffusionModelRunner.

The runner owns long-lived model-side state:

• the loaded pipeline and compilation setup;
• cache and offload integration;
• random-generator setup;
• request batches for complete forwards; and
• StepRequestState plus InputBatch for step execution.

This split keeps infrastructure out of model pipelines. A pipeline receives a request batch or step state and performs model work; it should not inspect engine queues or change scheduler state.

Lifecycle and cleanup

Startup

DiffusionEngine.make_engine() resolves the engine class, constructs the executor and scheduler, then normally runs a small dummy request to warm up the model.

Warmup is skipped only when distributed layerwise offload with AllGather is enabled (via enable_distributed_layerwise_offload and dlo_use_allgather) and max(data_parallel_size, sequence_parallel_size) > 1, because the dummy run sends one request while AllGather needs every shard rank to enter the same collective. If initialization or warmup fails, the runtime closes scheduler state and worker resources before returning the error.

Cancellation

abort() places request IDs on an abort queue. The busy loop marks those requests FINISHED_ABORTED; finalization removes scheduler state and emits an aborted output when a consumer still exists. The runner removes cached step-mode state when it receives a scheduler output that reports the finished request ID.

On the multiproc request-mode async path, abort or consumer drop while output is still materializing should release the associated async-output bookkeeping so late results are not retained after the request ends (pending #6253/#6439/#6580). Today, an unconsumed OUTPUT_READY is cached by design until wait_output_ready() or executor teardown. Dropping an output-stream consumer removes that consumer's queue. It does not take ownership of scheduler cleanup.

Shutdown and worker failure

DiffusionEngine.close() is idempotent. It stops the busy loop, wakes pending streams with an error, closes scheduler state, and shuts down the executor.

• mp: the executor asks workers to stop, waits for them, then terminates workers that miss the grace period. Shutdown marks unfinished RPC and async-output futures as failed, then clears them. Completed async outputs cached for later wait_output_ready() calls are dropped with the executor object.
• uni: the executor shuts down the in-process worker, drops the final model reference, and empties the accelerator cache so a later engine can reuse the device.

!!! warning

Do not continue using a worker group after a fatal collective timeout or
unexpected worker exit. Distributed state may be incomplete. The
multiprocess executor deliberately fails closed so the stage can restart.
The uniproc executor fails closed when the accelerator context itself is
poisoned.

Extension boundaries

• A custom engine may set default_diffusion_model_runner_cls; an explicit diffusion_model_runner_cls configuration still wins.
• Tests and custom engine integrations may inject a BaseScheduler subclass. SchedulerInterface remains only as a deprecated compatibility name.
• distributed_executor_backend may be "uni", "mp", a custom DiffusionExecutor subclass, or an import path. A backend must preserve scheduler output, request identity, health, and cleanup contracts. Use "mp" when you need process isolation or multi-GPU; "uni" is the single-GPU default.
• Worker extensions go through WorkerWrapperBase. They add worker methods; they do not become a second scheduler or request lifecycle.

These are advanced Python integration points, not stable end-user CLI customization surfaces.

Candidate invariants

DIFF-RUNTIME-INV-001: One lifecycle owner

Rule: Every admitted request MUST have exactly one scheduler-owned lifecycle until completion, cancellation, or failure.

DIFF-RUNTIME-INV-002: Execution follows scheduler output

Rule: Executors and workers MUST execute scheduler decisions without admitting, reordering, or forwarding requests independently. Distributed layerwise offload with AllGather is an exception at dispatch time: the scheduler may schedule a multi-request wave, but each DP rank executes one envelope from that wave so every rank enters the same collective schedule.

DIFF-RUNTIME-INV-003: Terminal cleanup is complete

Rule: Every terminal path MUST release request state, temporary tensors, hooks, and runtime-owned resources. Async-output bookkeeping on abort or consumer drop before consumption is also required (pending #6253/#6439/#6580).

DIFF-RUNTIME-INV-004: Optional features use runtime hooks

Rule: Cache, profiling, offload, and parallel features SHOULD integrate at defined hooks instead of creating another request lifecycle.

DIFF-RUNTIME-INV-005: Results keep request identity

Rule: Batched and asynchronous results MUST be mapped by stable request ID, not by assumed completion order.

Safe-change guide

Test the smallest affected slice, then cover its neighboring boundary:

Change area	Minimum evidence
Admission or state transitions	test_diffusion_scheduler.py, including duplicate IDs, compatibility, abort, and missing output
Engine loop or output delivery	test_diffusion_engine.py, test_diffusion_engine_cleanup.py, test_diffusion_engine_rpc_routing.py
Executor or IPC	test_multiproc_engine_concurrency.py, test_uniproc_executor.py, test_result_pump.py, test_diffusion_ipc.py, test_async_output_worker.py
Worker or runner	test_diffusion_worker.py, test_diffusion_model_runner.py
Stage boundary	test_stage_diffusion_proc.py, test_inline_stage_diffusion_client.py
Warmup or streaming	test_diffusion_engine_dummy_run.py, test_diffusion_streaming_output.py

Always exercise success, cancellation, per-request failure, fatal worker failure, repeated shutdown, one request, and multiple compatible requests. For step-mode changes, also test partial progress and cleanup of runner state.

Related documents

• Diffusion module overview
• Diffusion model integration
• Diffusion continuous batching
• Diffusion execution modes
• Async diffusion output
• Engine orchestration
• Stage runtime