Diffusion Model Runner for vLLM-Omni.
Handles model loading, compilation, caching, and execution of diffusion model forward passes. This follows the AR pattern where the Runner handles all model-related operations.
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).
instance-attribute ¶
diffusion_kv_backend = DiffusionKVModelRunnerBackend(
vllm_config=vllm_config,
od_config=od_config,
device=device,
)
instance-attribute ¶
kv_transfer_manager = (
payload_transfer_manager
if getattr(od_config, "kv_transfer_config", None)
is None
else None
)
Evict all cached text-encoder outputs (e.g. between training epochs).
Kept primarily for extension purposes.
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(
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(
scheduler_output: DiffusionSchedulerOutput,
) -> BatchRunnerOutput
Execute one step for one scheduled request and return runner output.
get_diffusion_kv_row(
request_id: str,
sequence_id: int | None,
context_id: str | None = None,
) -> int
Collect native specs from cache-enabled loaded attention modules.
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(
metadata: DiffusionKVMetadata,
) -> bool
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 | 'default' |
custom_pipeline_name | str | None | Optional custom pipeline class name to use. | None |
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.
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.
Physically initialize the Engine-generated rank-local config.
submit_interaction(
request_id: str, interaction: OmniInteractionPrompt
) -> None
Route a midway interaction through the pipeline interaction coordinator.