Skip to content

vllm_omni.diffusion.worker

Worker classes for diffusion models.

Modules:

Name Description
diffusion_model_runner

Diffusion Model Runner for vLLM-Omni.
diffusion_worker

Diffusion Worker for vLLM-Omni.
input_batch

Diffusion input-batch structures following the MRV2-style vLLM layout.
request_batch

Request-level batch abstraction for diffusion runner.
stage_payload
utils

Per-request mutable state for step-wise diffusion execution.

DiffusionModelRunner

Bases: DiffusionStagePayloadMixin

Model runner that handles model loading and execution for diffusion models.

This class follows the AR pattern where the Runner handles all model-related operations including loading, compilation, offloading, caching, and execution. The Worker only handles infrastructure (device, distributed env).

cache_backend instance-attribute 

cache_backend: Any | None = None

device instance-attribute 

device = device

diffusion_kv_backend instance-attribute 

diffusion_kv_backend = DiffusionKVModelRunnerBackend(
    vllm_config=vllm_config,
    od_config=od_config,
    device=device,
)

input_batch instance-attribute 

input_batch: InputBatch | None = None

kv_cache_config instance-attribute 

kv_cache_config: KVCacheConfig | None = None

kv_transfer_manager instance-attribute 

kv_transfer_manager = (
    payload_transfer_manager
    if getattr(od_config, "kv_transfer_config", None)
    is None
    else None
)

model_memory_usage instance-attribute 

model_memory_usage = 0

od_config instance-attribute 

od_config = od_config

offload_backend instance-attribute 

offload_backend: Any | None = None

pipeline instance-attribute 

pipeline: Any | None = None

prompt_embed_cache instance-attribute 

prompt_embed_cache: Any | None = None

state_cache instance-attribute 

state_cache: dict[str, StepRequestState] = {}

vllm_config instance-attribute 

vllm_config = vllm_config

clear_prompt_embed_cache 

clear_prompt_embed_cache() -> None

Evict all cached text-encoder outputs (e.g. between training epochs).

Kept primarily for extension purposes.

execute_model 

execute_model(
    req: OmniDiffusionRequest,
    kv_prefetch_job: KVPrefetchJob | None = None,
    diffusion_kv_metadata: DiffusionKVMetadata
    | None = None,
) -> DiffusionOutput

Execute a forward pass for the given requests.

Parameters:

Name Type Description Default
req OmniDiffusionRequest

A diffusion request containing a list of prompts to process.

 required

Returns:

Type Description
DiffusionOutput

DiffusionOutput with generated results.
Note

We use torch.no_grad() for HSDP because HSDP2's fully_shard requires access to tensor version counters in pre_forward hooks, which inference tensors do not track. For non-HSDP inference, we use torch.inference_mode() for better performance.

execute_model_batch 

execute_model_batch(
    scheduler_output: DiffusionSchedulerOutput,
    od_config: OmniDiffusionConfig,
) -> BatchRunnerOutput

Execute scheduled request-mode requests through the batch forward path.

Builds a DiffusionRequestBatch from scheduled new requests, runs per-request setup, and calls pipeline.forward(batch). The pipeline must declare supports_request_batch = True.

execute_stepwise 

execute_stepwise(
    scheduler_output: DiffusionSchedulerOutput,
) -> BatchRunnerOutput

Execute one step for one scheduled request and return runner output.

get_diffusion_kv_row 

get_diffusion_kv_row(
    request_id: str,
    sequence_id: int | None,
    context_id: str | None = None,
) -> int

get_kv_cache_spec 

get_kv_cache_spec() -> dict[str, KVCacheSpec]

Collect native specs from cache-enabled loaded attention modules.

get_prompt_embed_cache_stats 

get_prompt_embed_cache_stats() -> dict | None

Return hit/miss statistics for the prompt-embedding cache, if enabled.

Kept primarily for extension purposes.

install_diffusion_kv_metadata 

install_diffusion_kv_metadata(
    metadata: DiffusionKVMetadata,
) -> bool

load_model 

load_model(
    memory_pool_context_fn: Callable[
        [str], AbstractContextManager[Any]
    ]
    | None = None,
    load_format: str = "default",
    custom_pipeline_name: str | None = None,
) -> None

Launch the diffusion pipeline, applying compilation, offloading, and caching.

Parameters:

Name Type Description Default
memory_pool_context_fn Callable[[str], AbstractContextManager[Any]] | None

Optional function that returns a context manager for memory pool allocation (used for sleep mode).

 None
load_format str

Format for loading model weights. Supported formats: - "default" (default): Automatically detect and use the default format based on configuration - "custom_pipeline": Init model from a custom pipeline class specified by custom_pipeline_name - "dummy": Skip actual weight loading, useful for testing and custom pipelines that don't require default weights.

 'default'
custom_pipeline_name str | None

Optional custom pipeline class name to use.

 None

prepare_kv_for_forward 

prepare_kv_for_forward(
    scheduler_output: DiffusionSchedulerOutput,
)

profile_run 

profile_run(requests: list[OmniDiffusionRequest]) -> None

Run the maximum per-rank request batch for memory profiling.

This deliberately bypasses Scheduler admission and Diffusion KV metadata validation because cache capacity has not been sized yet. It otherwise uses the normal execution-mode path so model inputs, collective communication, backend workspaces, denoising, and decode allocations contribute to the observed peak. Step execution profiles one fused InputBatch instead of sequential single-request forwards.

refresh_diffusion_kv_block_table_layout 

refresh_diffusion_kv_block_table_layout() -> None

release_captured_graphs 

release_captured_graphs() -> None

Drop every CUDA graph held for this model, wherever it is kept.

Sleep level 2 discards the memory a capture was recorded against, so a graph that outlives it replays over freed storage. The runner owns compilation and execution resources, so it owns the release: a pipeline that captures graphs of its own implements release_captured_graphs and is collected here, instead of every caller having to know which pipelines have one.

remove_diffusion_kv_requests 

remove_diffusion_kv_requests(
    request_ids: Sequence[str | tuple[str, int]],
) -> int

set_kv_cache_config 

set_kv_cache_config(kv_cache_config: KVCacheConfig) -> None

Physically initialize the Engine-generated rank-local config.

submit_interaction 

submit_interaction(
    request_id: str, interaction: OmniInteractionPrompt
) -> None

Route a midway interaction through the pipeline interaction coordinator.

DiffusionWorker

A worker that manages GPU infrastructure and delegates to the model runner.

This class handles infrastructure initialization only: - Device setup (CUDA device selection) - Distributed environment (NCCL, model parallel) - Memory management (sleep/wake)

All model-related operations (loading, compilation, execution) are delegated to DiffusionModelRunner.

device instance-attribute 

device: device | None = None

init_snapshot instance-attribute 

init_snapshot: MemorySnapshot | None = None

local_rank instance-attribute 

local_rank = local_rank

lora_manager instance-attribute 

lora_manager: DiffusionLoRAManager | None = None

model_runner instance-attribute 

model_runner: DiffusionModelRunner | None = (
    model_runner_cls(
        vllm_config=self.vllm_config,
        od_config=self.od_config,
        device=self.device,
    )
)

od_config instance-attribute 

od_config = od_config

profiler instance-attribute 

profiler: WorkerProfiler | None = self._create_profiler()

rank instance-attribute 

rank = rank

requested_memory instance-attribute 

requested_memory: int | None = None

stage_id instance-attribute 

stage_id = getattr(od_config, 'stage_id', 0)

vllm_config instance-attribute 

vllm_config: VllmConfig | None = None

add_lora 

add_lora(lora_request: LoRARequest) -> bool

close_ar_diffusion_session 

close_ar_diffusion_session(session_id: str) -> bool

Close runner-owned AR state through the collective RPC boundary.

determine_available_kv_memory 

determine_available_kv_memory(
    profile_requests: list[OmniDiffusionRequest],
) -> list[int]

Profile and return each rank's safe Diffusion KV memory budget.

execute_model 

execute_model(
    req: OmniDiffusionRequest | list[NewRequestData],
    od_config: OmniDiffusionConfig,
    kv_prefetch_job: KVPrefetchJob | None = None,
    diffusion_kv_metadata: DiffusionKVMetadata
    | None = None,
) -> DiffusionOutput

Execute a forward pass by delegating to the model runner.

If req is a list (DP multi-concurrency), each rank picks one complete NewRequestData envelope based on its distributed rank. AllGather in the layerwise offload only gathers weight shards (request-independent), so all ranks stay synchronised at each AllGather call while computing different activations. Selecting the envelope keeps Scheduler-issued KV metadata bound to the request that owns its block tables.

Each rank returns its OWN DiffusionOutput (no gather). The executor collects N responses via the per-worker result queues.

execute_model_batch 

execute_model_batch(
    scheduler_output: DiffusionSchedulerOutput,
    od_config: OmniDiffusionConfig,
) -> BatchRunnerOutput

Batch forward: LoRA activate once, delegate to model runner.

execute_stepwise 

execute_stepwise(
    scheduler_output: DiffusionSchedulerOutput,
) -> BaseRunnerOutput

Execute one diffusion step by delegating to the model runner.

get_kv_cache_specs 

get_kv_cache_specs() -> list[dict[str, KVCacheSpec]]

Return native rank-local specs for every diffusion Worker.

handle_sleep_task 

handle_sleep_task(
    task: OmniSleepTask | dict,
) -> OmniACK | None

handle_wake_task 

handle_wake_task(
    task: OmniWakeTask | dict,
) -> OmniACK | None

init_device 

init_device() -> None

Initialize the device and distributed environment.

init_lora_manager 

init_lora_manager() -> None

Initialize the LoRA manager for this worker.

list_loras 

list_loras() -> list[int]

load_model 

load_model(
    load_format: str = "default",
    custom_pipeline_name: str | None = None,
    **kwargs,
) -> None

Load the diffusion model using DiffusionModelRunner.

pin_lora 

pin_lora(adapter_id: int) -> bool

prepare_kv_for_forward 

prepare_kv_for_forward(
    scheduler_output: DiffusionSchedulerOutput,
)

profile 

profile(
    is_start: bool = True, profile_prefix: str | None = None
) -> None

Start or stop profiling for this GPU worker.

Parameters:

Name Type Description Default
is_start bool

True to start profiling, False to stop.

 True
profile_prefix str | None

Optional prefix for trace filename.

 None

remove_diffusion_kv_requests 

remove_diffusion_kv_requests(
    request_ids: list[str | tuple[str, int]],
) -> int

Clear Worker-local rows without freeing Scheduler-owned blocks.

remove_lora 

remove_lora(adapter_id: int) -> bool

reset_ar_diffusion_session 

reset_ar_diffusion_session(session_id: str) -> bool

Reset runner-owned AR state through the collective RPC boundary.

set_kv_cache_configs 

set_kv_cache_configs(
    kv_cache_configs: list[KVCacheConfig],
    resolved_max_model_len: int,
) -> None

Select this rank's config and initialize its physical KV pages.

shutdown 

shutdown() -> None

Shutdown the worker and release process-global resources.

sleep 

sleep(level: int = 1) -> int

Put the worker to sleep, offloading model weights.

Parameters:

Name Type Description Default
level int

Sleep level. Level 1 offloads weights, level 2 also saves buffers.

 1

submit_interaction 

submit_interaction(
    request_id: str, interaction: OmniInteractionPrompt
) -> None

Apply a midway interaction to an active stepwise request.

synchronize_device 

synchronize_device(timeout: float | None = None) -> None

Wait until this rank has no device work left.

A KV prefetch still receiving on its background thread has queued no device work yet, so it is joined first. timeout bounds that join (and the multi-process worker's output drain before this call); the device wait itself is unbounded.

wake_up 

wake_up(tags: list[str] | None = None) -> bool

Wake up the worker from sleep mode.

Re-activates the memory allocator for the specified tags and restores model buffers from CPU back to GPU if they were saved during Level 2 sleep.

Parameters:

Name Type Description Default
tags list[str] | None

List of memory pool tags to re-activate (e.g., ["weights"] to match Level 1 sleep). If None, all pools are re-activated.

 None

WorkerProc

Wrapper that runs one Worker in a separate process.

context instance-attribute 

context = zmq.Context(io_threads=2)

gpu_id instance-attribute 

gpu_id = gpu_id

mq instance-attribute 

mq = MessageQueue.create_from_handle(
    broadcast_handle, gpu_id
)

od_config instance-attribute 

od_config = od_config

result_mq instance-attribute 

result_mq = MessageQueue(
    n_reader=1, n_local_reader=1, local_reader_ranks=[0]
)

result_mq_handle instance-attribute 

result_mq_handle = self.result_mq.export_handle()

wake_event instance-attribute 

wake_event = wake_event

worker instance-attribute 

worker = self._create_worker(
    gpu_id,
    od_config,
    worker_extension_cls,
    custom_pipeline_args,
)

drain_async_outputs 

drain_async_outputs(timeout: float | None = None) -> bool

Block until background D2H/SHM packing has no work left.

Returns False if outputs are still in flight when timeout expires.

recv_message 

recv_message() -> Any

Receive one complete broadcast message without dropping overflow data.

shutdown 

shutdown() -> None

Stop background work and release worker-owned IPC resources.

worker_main staticmethod 

worker_main(
    rank: int,
    od_config: OmniDiffusionConfig,
    pipe_writer: Connection,
    broadcast_handle,
    wake_event: Event,
    worker_extension_cls: str | None = None,
    custom_pipeline_args: dict[str, Any] | None = None,
) -> None

Worker initialization and execution loops.