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

---
title: Diffusion Runtime
kind: module
status: draft
owners:
  - "@Isotr0py"
  - "@princepride"
  - "@SamitHuang"
  - "@fhfuih"
primary_code_paths:
  - 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/**
related_code_paths:
  - 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/**
depends_on:
  - ../engine_orchestration.md
  - ../stage_runtime.md
  - ../input_output_modality_contracts.md
  - diffusion_model_integration.md
validation_paths:
  - 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_ray_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
upstream_refs:
  - diffusers.DiffusionPipeline
last_reviewed: 2026-09-02
last_verified_commit: e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6
---

# 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](../engine_orchestration.md);
- process placement and replica lifecycle belong to
  [Stage runtime](../stage_runtime.md); and
- model-specific work belongs to
  [Diffusion model integration](diffusion_model_integration.md).

## 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.
- **`ray`** (explicit): `RayDiffusionExecutor` starts one GPU actor per worker
  across one or more Ray nodes.

In all backends, each worker owns its device, distributed state,
`DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance.

```mermaid
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, mp, or ray 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
    executor -->|"ray: actor RPC"| rayActor["RayDiffusionWorkerWrapper"]
    rayActor --> 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 |
| `RayDiffusionExecutor` | GPU actors, placement groups, remote RPC, actor monitoring | Local worker subprocesses and shared-memory IPC |
| `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.

```mermaid
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 dispatch path supported by `mp` and `ray`. 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`; these executors route 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, both backends reject 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:

```mermaid
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` | Set explicitly | One GPU actor per worker across Ray nodes |

The external-launcher diffusion backend is 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](../../feature/async_diffusion_output.md) for the
multiproc request-mode timeline.

### Ray backend

`RayDiffusionExecutor` reuses initialized Ray or calls `ray.init()` (respecting
`RAY_ADDRESS`). It reuses the current placement group or reserves one GPU
bundle per worker with `PACK` placement. Ray owns GPU placement; PyTorch
owns distributed collectives through a shared TCP rendezvous. Global ranks
are grouped by host, while each actor uses local device 0.

Request batches, step execution, and DLO DP waves use actor RPCs with
rank-aware replies. Returned tensors move to CPU before Ray serialization;
this backend does not use the multiproc async output pumps. Driver inference
settings and explicit stage `runtime.env` are forwarded, excluding worker
identity and device assignment. See [Stage configuration](../../../configuration/stage_configs.md)
for cluster setup and environment details.

## 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. Ray actors also construct `WorkerWrapperBase`. In all 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.
- **`ray`**: the executor gives actors up to five seconds to shut down, then
  kills them and removes placement groups it owns. Actor death is monitored
  even while idle; actor errors and RPC timeouts fail the group and notify the
  stage. The stage client allows subprocess cleanup before signal fallback.
- **`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
    `mp` and `ray` executors deliberately fail 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"`, `"ray"`, 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
  local process isolation or multi-GPU, or `"ray"` for workers across nodes;
  `"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_ray_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](index.md)
- [Diffusion model integration](diffusion_model_integration.md)
- [Diffusion continuous batching](continuous_batching.md)
- [Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md)
- [Async diffusion output](../../feature/async_diffusion_output.md)
- [Engine orchestration](../engine_orchestration.md)
- [Stage runtime](../stage_runtime.md)