Diffusion Module Architecture Design¶

The vLLM-Omni diffusion module (vllm_omni/diffusion) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.

This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.

Main Components of the Diffusion Module

Table of Content:

Architecture Overview
Diffusion Engine
Scheduler
Worker
Diffusion Pipeline
Acceleration Components
Data Flow

Architecture Overview¶

The diffusion module follows a multi-process, distributed architecture with clear separation of concerns:

Diffusion Architecture Overview

1. Diffusion Engine¶

Location: vllm_omni/diffusion/diffusion_engine.py

Responsibilities¶

The DiffusionEngine is the orchestrator of the diffusion inference system. It owns request admission, scheduler/executor coordination, output-stream delivery, cancellation cleanup, and startup warmup.

Key Components¶

1.1 Initialization¶

class DiffusionEngine:
    def __init__(
        self,
        od_config: OmniDiffusionConfig,
        scheduler: SchedulerInterface | None = None,
    ):
        self.od_config = od_config
        self._init_process_hooks(od_config)
        self.execution_mode = self._resolve_execution_mode(od_config)
        self._init_executor(od_config)
        self._init_scheduler(od_config, scheduler)
        self._init_runtime_state()
        self._init_execute_fn()

Key Features:

Pre/Post Processing: Registers model-specific pre-processing and post-processing functions via registry pattern
Execution Mode Selection: Selects one of two execution modes:
REQUEST_BATCH: complete requests are scheduled through request-level execution. max_num_seqs=1 is the serial request path; values above 1 allow compatible requests to use fused pipeline.forward(batch) when the pipeline supports it.
STEP_BATCH: requests are advanced one denoising step at a time through prepare_encode(), denoise_step(), step_scheduler(), and post_decode(). max_num_seqs controls the maximum number of compatible active requests in one step wave.
Scheduler Selection: Uses StepScheduler for step-batch execution and RequestScheduler for request-batch execution. Tests and custom integrations may inject a SchedulerInterface instance explicitly.
Executor Selection: Builds the configured DiffusionExecutor, then routes scheduled work through execute_batch or execute_step according to the execution mode. execute_batch handles the serial max_num_seqs=1 request path internally.
Factory-Owned Warmup: DiffusionEngine.make_engine() constructs the selected engine backend and then runs startup warmup through run_startup_warmup().

1.2 Execution Modes¶

class DiffusionExecutionMode(str, Enum):
    REQUEST_BATCH = "request_batch"
    STEP_BATCH = "step_batch"

streaming_output=True requires step execution. If streaming is enabled without step_execution=True, the engine enables step execution so the output stream can receive intermediate chunks or the final step output.

1.3 Output Stream Lifecycle¶

request_id = engine.add_request(request)
async for output in engine.get_output_stream(request_id):
    yield output

Design Decisions:

Unified Stream Semantics: step_streaming() is the public async engine path. Streaming callers forward every yielded output. Non-streaming callers consume the same stream to completion and return only the final output.
Queue Delivery: Each admitted request gets an asyncio.Queue[DiffusionOutput]. The busy loop emits chunks/final outputs into that queue.
Cancellation Cleanup: If a client disconnects or cancels stream consumption, get_output_stream() removes only the queue. Scheduler terminal state is still finalized when the scheduler reports the request as finished, so request IDs do not remain active indefinitely.

1.4 Request Processing Flow¶

async def step_streaming(self, request: OmniDiffusionRequest):
    # 1. Pre-process requests
    request = self.pre_process_func(request)

    # 2. Add request and consume the unified output stream
    async for output in self.async_add_req_and_stream_response(request):
        # 3. Post-process each delivered output
        yield self.postprocess_output(request, output, ...)

Flow:

Pre-processing: Applies model-specific transformations
Scheduling: Delegates request state transitions to the selected scheduler
Execution: Calls the selected executor function for request, request-batch, or step execution
Output Delivery: Emits intermediate chunks and final outputs through the request output stream
Post-processing: Converts raw outputs to final format (e.g., PIL images)

2. Scheduler¶

Location: vllm_omni/diffusion/sched/

Architecture¶

The scheduler is a request-state scheduler. It owns request lifecycle management and scheduling decisions, while execution stays in DiffusionEngine and the executor.

Key Components¶

2.1 Request State Model¶

class DiffusionRequestStatus(enum.IntEnum):
    WAITING = ...
    RUNNING = ...
    PREEMPTED = ...
    FINISHED_COMPLETED = ...
    FINISHED_ABORTED = ...
    FINISHED_ERROR = ...

@dataclass
class SchedulerRequestState:
    request_id: str
    req: OmniDiffusionRequest
    sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
    status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
    error: str | None = None

Design Features:

Explicit lifecycle: Requests move through waiting, running, optional preemption, and terminal states.
Centralized error handling: Completion, abort, and error states are all normalized in the scheduler layer.

2.2 Shared Bookkeeping in `BaseScheduler`¶

class BaseScheduler:
    def __init__(self) -> None:
        self._request_states = {}
        self._waiting = deque()
        self._running = []
        self._finished_req_ids = set()
        self.max_num_running_reqs = 1

    def add_request(self, request: OmniDiffusionRequest) -> str: ...
    def schedule(self) -> DiffusionSchedulerOutput: ...
    def update_from_output(
        self,
        sched_output: DiffusionSchedulerOutput,
        output: BaseRunnerOutput,
    ) -> set[str]: ...

Responsibilities:

Lifecycle contract: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
Stable boundary: DiffusionSchedulerOutput is the only scheduling result consumed by DiffusionEngine.
Pluggability: Different scheduler policies can reuse the same engine integration path.

Design Features:

Common state storage: Shared request maps and waiting/running sets live in the base class.
Shared cleanup logic: Duplicate-request checks, finish handling, and state removal are centralized instead of duplicated in each policy.
Current constraint boundary: BaseScheduler derives max_num_running_reqs from max_num_seqs. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares supports_request_batch = True; step-wise diffusion uses the same scheduler capacity for compatible step batches.

2.3 Current `RequestScheduler` Policy¶

class RequestScheduler(BaseScheduler):
    def schedule(self) -> DiffusionSchedulerOutput:
        # 1. keep existing RUNNING requests in the scheduling result
        # 2. pull WAITING requests while capacity remains
        # 3. move newly admitted requests into RUNNING

Behavior:

FIFO request scheduling: Waiting requests are promoted in queue order.
Compatible request admission: RequestScheduler admits waiting requests while capacity remains and the request's RequestBatchSamplingParamsKey is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side DiffusionRequestBatch.
Executor result feedback: update_from_output() converts executor output into FINISHED_COMPLETED or FINISHED_ERROR and returns finished request ids.

2.4 Engine-Driven Execution Loop¶

request_id = scheduler.add_request(request)
while True:
    sched_output = scheduler.schedule()
    output = executor.add_req(req)
    finished_req_ids = scheduler.update_from_output(sched_output, output)

Design Decisions:

Separation of concerns: Scheduler manages state and policy; executor handles runtime execution.
No scheduler-owned IPC: Scheduler no longer talks to workers directly.
Split concurrency model: Request-mode diffusion can schedule a static batch of compatible independent requests for one full pipeline forward when the pipeline supports request-level batching; the step-wise path can also admit or remove compatible requests between denoise steps. See Diffusion Continuous Batching.

3. Worker¶

Location: vllm_omni/diffusion/worker/gpu_worker.py

Architecture¶

Workers are independent processes that execute the actual model inference. Each worker runs on a dedicated GPU and participates in distributed inference.

Key Components¶

3.1 Worker Process Structure¶

class WorkerProc:
    def __init__(self, od_config, gpu_id, broadcast_handle):
        # Initialize ZMQ context for IPC
        self.context = zmq.Context(io_threads=2)

        # Connect to broadcast queue (receive requests)
        self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)

        # Create result queue (only rank 0)
        if gpu_id == 0:
            self.result_mq = MessageQueue(n_reader=1, ...)

        # Initialize GPU worker
        self.worker = self._create_worker(gpu_id, od_config, ...)

Initialization Steps:

IPC Setup: Creates ZMQ context and message queues
Distributed Environment Setup: Initializes PyTorch distributed communication
- For CUDA GPUs: Uses NCCL (fast GPU communication)
- For NPU: Uses HCCL (Huawei Collective Communications Library)
- For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
Model Loading: Loads diffusion pipeline on assigned GPU
Cache Setup: Enables cache backend if configured.

3.2 GPU Worker¶

class DiffusionWorker:
    def init_device(self):
        # Set distributed environment variables
        os.environ["RANK"] = str(rank)
        os.environ["WORLD_SIZE"] = str(world_size)

        # Select the device and initialize distributed execution
        self.device = current_omni_platform.get_torch_device(rank)
        current_omni_platform.set_device(self.device)
        self.vllm_config = _create_diffusion_worker_vllm_config(
            self.device, self.od_config
        )
        init_distributed_environment(world_size, rank)
        parallel_config = self.od_config.parallel_config
        initialize_model_parallel(
            data_parallel_size=parallel_config.data_parallel_size,
            cfg_parallel_size=parallel_config.cfg_parallel_size,
            sequence_parallel_size=parallel_config.sequence_parallel_size,
            tensor_parallel_size=parallel_config.tensor_parallel_size,
            pipeline_parallel_size=parallel_config.pipeline_parallel_size,
        )

    def load_model(self, load_format="default", custom_pipeline_name=None):
        self.model_runner.load_model(
            load_format=load_format,
            custom_pipeline_name=custom_pipeline_name,
        )

    def init_lora_manager(self):
        self.lora_manager = DiffusionLoRAManager(
            pipeline=self.model_runner.pipeline,
            device=self.device,
            dtype=self.od_config.dtype,
        )

Key Features:

Tensor Parallelism: Supports multi-GPU tensor parallelism via PyTorch distributed
Model Loading: Uses DiffusersPipelineLoader for efficient weight loading
Cache Integration: Enables cache backends (TeaCache, cache-dit, etc.) transparently

3.3 Worker Busy Loop¶

def _worker_busy_loop(self):
    while self._running:
        # 1. Receive unified message (generation request, RPC request, or shutdown)
        msg = self.mq.dequeue(timeout=1.0)

        # 2. Route message based on type
        if isinstance(msg, dict) and msg.get("type") == "rpc":
            # Handle RPC request
            result, should_reply = self._execute_rpc(msg)
            if should_reply:
                self._return_result(result)

        elif isinstance(msg, dict) and msg.get("type") == "shutdown":
            # Handle shutdown message
            self._running = False

        else:
            # Handle generation request (OmniDiffusionRequest list)
            output = self.worker.execute_model(msg, self.od_config)
            self._return_result(output)

Execution Flow:

Receive: Dequeues unified messages from shared memory queue
Route: Handles different message types (generation, RPC, shutdown)
Execute: Runs forward pass through pipeline for generation requests
Respond: Sends results back (rank 0 for generation, specified rank for RPC)

3.4 Model Execution¶

@torch.inference_mode()
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config):
    req = reqs[0]  # TODO: support batching

    # Refresh cache backend if enabled
    if self.cache_backend is not None and self.cache_backend.is_enabled():
        self.cache_backend.refresh(self.pipeline, req.num_inference_steps)

    # Set forward context for parallelism
    with set_forward_context(
        vllm_config=self.vllm_config,
        omni_diffusion_config=self.od_config
    ):
        output = self.pipeline.forward(req)
    return output

The model execution leverages multiple parallelism strategies that are transparently applied during the forward pass. The set_forward_context() context manager makes parallel group information available throughout the forward pass:

# Inside transformer layers, parallel groups are accessed via:
from vllm_omni.diffusion.distributed.parallel_state import (
    get_sp_group, get_dp_group, get_cfg_group, get_pp_group
)

Optimizations:

Cache Refresh: Clears cache state before each generation for clean state
Context Management: Forward context ensures parallel groups are available during execution
Single Request: Currently processes one request at a time (batching TODO)

4. Diffusion Pipeline¶

Location: vllm_omni/diffusion/models/*/pipeline_*.py

The pipeline is the model-specific implementation that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.

Most pipeline implementation are referred from diffusers. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the diffuse function in the pipeline class. An example is as follows:

def diffuse(self, ...):
    for i, t in enumerate(timesteps):
        # Forward pass for positive prompt
        transformer_kwargs = {
            "hidden_states": latents,
            "timestep": timestep / 1000,
            "encoder_hidden_states": prompt_embeds,
        }
        noise_pred = self.transformer(**transformer_kwargs)[0]

        # Forward pass for negative prompt (CFG)
        if do_true_cfg:
            neg_transformer_kwargs = {...}
            neg_transformer_kwargs["cache_branch"] = "negative"
            neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]

            # Combine predictions
            comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
            noise_pred = comb_pred * (cond_norm / noise_norm)

        # Scheduler step
        latents = self.scheduler.step(noise_pred, t, latents)[0]

    return latents

Key Features:

CFG Support: Handles classifier-free guidance with separate forward passes
Cache Branching: Uses cache_branch parameter for cache-aware execution
True CFG: Implements advanced CFG with norm preservation

To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view Adding Diffusion Model

5. Acceleration Components¶

5.1 Attention Backends¶

Location: vllm_omni/diffusion/attention/

Architecture¶

The attention system uses a role-aware backend selector pattern. Each Attention site declares a semantic role ("self", "cross", or a model-specific string). The selector consults the user's AttentionConfig to pick a backend per role and falls back to the platform default when nothing is configured.

Backend Selection¶

Location: vllm_omni/diffusion/attention/selector.py

class Attention(nn.Module):
    def __init__(self, num_heads, head_size, causal, softmax_scale, *,
                 role="self", role_category=None, ...):
        # Resolve backend for this role from the active OmniDiffusionConfig
        config = get_current_diffusion_config_or_none()
        attention_config = config.diffusion_attention_config if config is not None else None
        attn_backend_cls, spec = get_attn_backend_for_role(
            role=role,
            head_size=head_size,
            attention_config=attention_config,
            role_category=role_category,
        )
        self.attn_impl_cls = attn_backend_cls.get_impl_cls()
        self.attention = self.attn_impl_cls(..., backend_kwargs=spec.backend_kwargs() if spec else None)

Available Backends:

FlashAttention (FLASH_ATTN): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and mindiesd on Ascend NPU.
SDPA (TORCH_SDPA): PyTorch's scaled dot-product attention — cross-platform fallback.
SageAttention (SAGE_ATTN): Quantized attention implementation from the SageAttention library.

These backends provide the kernel implementations for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see Parallel Attention.

Backend Selection Mechanism¶

Selection is driven by OmniDiffusionConfig.diffusion_attention_config, which carries a global default plus a per_role map of AttentionSpec entries. A role string is resolved in this order:

def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
    # 1. attention_config.per_role[role]            — exact match
    # 2. attention_config.per_role[role_category]   — category fallback
    # 3. attention_config.default                   — global default
    # 4. Platform default                           — hardware-specific
    if attention_config is not None:
        spec, source = attention_config.resolve_with_source(
            role=role, role_category=role_category,
        )
        if spec is not None:
            return load_backend(spec.backend), spec
    return current_omni_platform.get_diffusion_attn_backend_cls(...), None

Selection Priority:

Per-role override (--diffusion-attention-config.per_role.<role>.backend) — finest control; matched on the layer's exact role string.
Role category fallback (--diffusion-attention-config.per_role.<category>.backend) — used when an exact match is missing and the layer declared role_category (e.g. a "mymodel.audio_to_video" site can fall back to "cross").
Global default (--diffusion-attention-backend FLASH_ATTN, or --diffusion-attention-config.default.backend, or DIFFUSION_ATTENTION_BACKEND env var).
Platform default — current_omni_platform.get_diffusion_attn_backend_cls(...) picks the best-available kernel for the hardware.

A backend that needs configuration exposes it as a typed field on AttentionSpec; AttentionSpec.backend_kwargs() serializes it into the backend_kwargs the impl consumes. Today those fields are skip_softmax (TRTLLM_ATTN) and block_sparse (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. target_sparsity or sparsity can be set without changing model code.

For the user-facing CLI surface, see Diffusion Attention Backends. For the role declaration contract on the model side, see Adding a Diffusion Model.

Backend Availability:

SDPA: Always available (PyTorch built-in)
FlashAttention: Requires flash-attn on CUDA, aiter on ROCm, or mindiesd on Ascend NPU
SageAttention: Requires sage-attention package (from THU-ML GitHub)

Attention Backend Registry¶

Location: vllm_omni/diffusion/attention/selector.py + vllm_omni/platforms/<device>/platform.py

Backend resolution is delegated to the active platform: current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size) returns a fully-qualified class path that _load_backend_cls imports lazily, with the result cached by (backend_name, head_size).

@cache
def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]:
    backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
        selected_backend=backend_name,
        head_size=head_size,
    )
    return _load_backend_cls(backend_cls_path)

Each platform (cuda, rocm, xpu, musa, npu) maps backend names like "FLASH_ATTN", "TORCH_SDPA", "SAGE_ATTN" to the right backend class path for that hardware, including head-size compatibility checks. Passing selected_backend=None lets the platform pick its own default.

Attention Backend Integration¶

The Attention layer integrates backends through a unified interface. Here's how FlashAttentionBackend is integrated as an example:

# attention/backends/flash_attn.py

class FlashAttentionBackend(AttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "FLASH_ATTN"

    @staticmethod
    def get_impl_cls() -> type["FlashAttentionImpl"]:
        return FlashAttentionImpl

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return [64, 96, 128, 192, 256]  # FlashAttention supports these head sizes


class FlashAttentionImpl(AttentionImpl):
    def __init__(self, num_heads, head_size, softmax_scale, causal, ...):
        self.num_heads = num_heads
        self.causal = causal
        self.softmax_scale = softmax_scale

    def forward(self, query, key, value, attn_metadata=None):
        # Call FlashAttention kernel
        out = flash_attn_func(
            query, key, value,
            causal=self.causal,
            softmax_scale=self.softmax_scale,
        )
        return out

5.2 Parallel Attention¶

Location: vllm_omni/diffusion/attention/parallel/

Architecture¶

Parallel attention strategies implement Sequence Parallelism (SP) at the attention layer level. These strategies distribute attention computation across multiple GPUs by splitting the sequence dimension, using different communication patterns. They work on top of AttentionBackend implementations (FlashAttention, SDPA, etc.), handling the parallelization/communication while the backends handle the actual attention computation.

Key Distinction: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the ParallelAttentionStrategy interface and use AttentionBackend implementations internally.

Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:

Split the sequence dimension across GPUs
Contribute to sequence_parallel_size (via ring_degree and ulysses_degree)
Work at the attention layer level (not model/pipeline level)

Ulysses Sequence Parallelism (USP)¶

Location: vllm_omni/diffusion/attention/parallel/ulysses.py

USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses all-to-all communication to efficiently parallelize attention for very long sequences. Specifically, it uses all-to-all collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.

Ulysses splits attention computation in two dimensions:

Sequence Dimension: Splits the sequence length across GPUs
Head Dimension: Splits attention heads across GPUs

Configuration: ulysses_degree contributes to sequence_parallel_size

Ring Sequence Parallelism¶

Location: vllm_omni/diffusion/attention/parallel/ring.py

Ring Attention is a parallel attention strategy that implements sequence parallelism using ring-based point-to-point (P2P) communication. Unlike attention backends that provide the attention kernel implementation, Ring Attention is a communication pattern that works on top of attention backends (FlashAttention or SDPA).

Ring Attention splits sequence dimension across GPUs in a ring topology, implemented via the ParallelAttentionStrategy interface, instead of AttentionBackend. P2P ring communication is applied to circulate Key/Value blocks across GPUs. Internally, ring_flash_attn_func or ring_pytorch_attn_func is used depending on available backends.

Architecture:

class RingParallelAttention:
    """Ring sequence-parallel strategy."""

    def run_attention(self, query, key, value, attn_metadata, ...):
        # Selects underlying attention kernel (FlashAttention or SDPA)
        if backend_pref == "sdpa":
            return ring_pytorch_attn_func(...)  # Uses SDPA kernel
        else:
            return ring_flash_attn_func(...)    # Uses FlashAttention kernel

Integration:

Ring Attention is activated when ring_degree > 1 in parallel config
It's selected by build_parallel_attention_strategy() in the attention layer
The Attention layer routes to _run_ring_attention() when Ring is enabled
Works alongside attention backends: Ring handles communication, backends handle computation

Configuration: ring_degree contributes to sequence_parallel_size

Relationship with AttentionBackend¶

Parallel attention strategies (Ring, Ulysses) work on top of AttentionBackend implementations:

They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
They handle the multi-GPU communication/parallelization layer
They implement ParallelAttentionStrategy interface (not AttentionBackend)

For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see Parallel Strategies.

5.3 Cache Backends¶

Location: vllm_omni/diffusion/cache/

Architecture¶

Cache backends provide a unified interface for applying different caching strategies to accelerate diffusion inference. The system supports multiple backends (TeaCache, cache-dit) with a consistent API for enabling and refreshing cache state.

Cache Backend Interface¶

class CacheBackend(ABC):
    def __init__(self, config: DiffusionCacheConfig):
        self.config = config
        self.enabled = False

    @abstractmethod
    def enable(self, pipeline: Any) -> None:
        """Enable cache on the pipeline."""
        raise NotImplementedError

    @abstractmethod
    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
        """Refresh cache state for new generation."""
        raise NotImplementedError

    def is_enabled(self) -> bool:
        """Check if cache is enabled."""
        return self.enabled

Design Pattern:

Abstract Base Class: Defines contract for all cache backends
Pipeline-based: Works with pipeline instances (not just transformers)
State Management: Provides refresh mechanism for clean state between generations

Available Backends¶

1. TeaCache Backend

Location: vllm_omni/diffusion/cache/teacache/backend.py

class TeaCacheBackend(CacheBackend):
    def enable(self, pipeline: Any):
        # Extract transformer from pipeline
        transformer = pipeline.transformer
        transformer_type = transformer.__class__.__name__

        # Create TeaCacheConfig from DiffusionCacheConfig
        teacache_config = TeaCacheConfig(
            transformer_type=transformer_type,
            rel_l1_thresh=self.config.rel_l1_thresh,
            coefficients=self.config.coefficients,
        )

        # Apply hooks to transformer
        apply_teacache_hook(transformer, teacache_config)
        self.enabled = True

    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
        transformer = pipeline.transformer
        if hasattr(transformer, "_hook_registry"):
            transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)

TeaCache Features:

Timestep-aware: Caches based on timestep embedding similarity
Adaptive: Dynamically decides when to reuse cached computations
CFG-aware: Handles positive/negative branches separately
Custom Hook System: Uses a custom forward interception mechanism (via HookRegistry) that wraps the module's forward method, allowing transparent integration without modifying model code

2. Cache-DiT Backend

Location: vllm_omni/diffusion/cache/cachedit/

class CacheDiTBackend(CacheBackend):
    def enable(self, pipeline: Any):
        # Uses cache-dit library for acceleration
        # Supports DBCache, SCM (Step Computation Masking), TaylorSeer
        # Works with single and dual-transformer architectures
        ...
        self.enabled = True

    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
        # Updates cache context with new num_inference_steps
        ...

Cache-DiT Features:

DBCache: Dynamic block caching with configurable compute blocks
SCM: Step Computation Masking for additional speedup
TaylorSeer: Advanced calibration for cache accuracy
Dual-transformer Support: Handles models like Wan2.2 with two transformers

Cache Backend Selector¶

Location: vllm_omni/diffusion/cache/selector.py

def get_cache_backend(
    cache_backend: str | None,
    cache_config: dict | DiffusionCacheConfig
) -> CacheBackend | None:
    """Get cache backend instance based on cache_backend string.

    Args:
        cache_backend: Cache backend name ("cache_dit", "tea_cache", or None)
        cache_config: Cache configuration (dict or DiffusionCacheConfig)

    Returns:
        Cache backend instance or None if cache_backend is "none"
    """
    if cache_backend is None or cache_backend == "none":
        return None

    if isinstance(cache_config, dict):
        cache_config = DiffusionCacheConfig.from_dict(cache_config)

    if cache_backend == "cache_dit":
        return CacheDiTBackend(cache_config)
    elif cache_backend == "tea_cache":
        return TeaCacheBackend(cache_config)
    else:
        raise ValueError(f"Unsupported cache backend: {cache_backend}")

Usage Flow:

Selection: get_cache_backend() returns appropriate backend instance
Enable: backend.enable(pipeline) called during worker initialization
Refresh: backend.refresh(pipeline, num_inference_steps) called before each generation
Check: backend.is_enabled() verifies cache is active

5.4 Parallel Strategies¶

Location: vllm_omni/diffusion/distributed/parallel_state.py

Parallelism Types¶

The system supports multiple orthogonal parallelism strategies:

Sequence Parallelism (SP)

Purpose: Split sequence dimension across GPUs
Attention-level SP: Ring Attention and Ulysses (USP) implement SP at the attention layer level
- See Parallel Attention for details
- Configuration: ulysses_degree × ring_degree = sequence_parallel_size
Use Case: Very long sequences (e.g., high-resolution images)

Data Parallelism (DP)

Purpose: Replicate model across GPUs, split batch
Use Case: Batch processing, throughput optimization

Tensor Parallelism (TP) (Experimental)

Purpose: Split model weights across GPUs
Implementation: Uses vLLM's tensor parallel groups
Use Case: Large models that don't fit on single GPU

CFG Parallelism (under development)

Purpose: Parallelize Classifier-Free Guidance (positive/negative prompts)
Infrastructure: CFG parallel groups are initialized and available via get_cfg_group()

Parallel Group Management¶

def initialize_model_parallel(
    data_parallel_size: int = 1,
    cfg_parallel_size: int = 1,
    sequence_parallel_size: int | None = None,
    ulysses_degree: int = 1,
    ring_degree: int = 1,
    tensor_parallel_size: int = 1,
    pipeline_parallel_size: int = 1,
    vae_parallel_size: int = 0,
):
    # Generate orthogonal parallel groups
    rank_generator = RankGenerator(
        tensor_parallel_size,
        sequence_parallel_size,
        pipeline_parallel_size,
        cfg_parallel_size,
        data_parallel_size,
        "tp-sp-pp-cfg-dp",
    )

    # Initialize each parallel group
    _DP = init_model_parallel_group(rank_generator.get_ranks("dp"), ...)
    _CFG = init_model_parallel_group(rank_generator.get_ranks("cfg"), ...)
    _SP = init_model_parallel_group(rank_generator.get_ranks("sp"), ...)
    _PP = init_model_parallel_group(rank_generator.get_ranks("pp"), ...)
    _TP = init_model_parallel_group(rank_generator.get_ranks("tp"), ...)

Rank Order: tp-sp-pp-cfg-dp (tensor → sequence → pipeline → cfg → data)

Note: For attention-level Sequence Parallelism implementations (Ring Attention and Ulysses), see Parallel Attention. This section covers higher-level parallelism strategies.

6. Data Flow¶

Complete Request Flow¶

End-to-end Data Flow in the vLLM-Omni Diffusion Module

1. User Request
   └─> Omni.generate(prompt)
       └─> Prepare OmniDiffusionRequest
           └─> DiffusionEngine.step_streaming(request)

2. Pre-processing
   └─> pre_process_func(request)
       └─> Model-specific transformations

3. Scheduling
       └─> scheduler.add_request(request)
           └─> scheduler.schedule()
           └─> DiffusionEngine selects execute_batch / execute_step

4. Worker Execution
   └─> WorkerProc._worker_busy_loop()
       └─> DiffusionWorker.execute_model(req)
           └─> DiffusionModelRunner selects execution mode
               ├─> REQUEST_BATCH
               │   └─> Pipeline.forward(req_or_batch)
               │       ├─> encode_prompt()
               │       ├─> prepare_latents()
               │       ├─> diffuse() [loop]
               │       │   ├─> transformer.forward() [with cache backend hooks]
               │       │   └─> scheduler.step()
               │       └─> vae.decode()
               └─> STEP
                   ├─> pipeline.prepare_encode(state) [new requests]
                   ├─> pipeline.denoise_step(input_batch)
                   ├─> pipeline.step_scheduler(state, noise_pred)
                   └─> pipeline.post_decode(state) [chunk boundary or final step]

5. Result Collection
   └─> Executor returns RunnerOutput / BatchRunnerOutput
       └─> scheduler.update_from_output(...)
           ├─> DiffusionEngine emits chunk/final output to request queue
           └─> DiffusionEngine finalizes terminal scheduler state

6. Post-processing
   └─> post_process_func(output)
       └─> Convert to PIL images / final format

7. Stream Completion / Cancellation
   ├─> Streaming callers receive every yielded output
   ├─> Non-streaming callers drain the stream and return the final output
   └─> Cancelled streams drop their output queue but still allow scheduler state cleanup

Diffusion Module Architecture Design¶

Architecture Overview¶

1. Diffusion Engine¶

Responsibilities¶

Key Components¶

1.1 Initialization¶

1.2 Execution Modes¶

1.3 Output Stream Lifecycle¶

1.4 Request Processing Flow¶

2. Scheduler¶

Architecture¶

Key Components¶

2.1 Request State Model¶

2.2 Shared Bookkeeping in BaseScheduler¶

2.3 Current RequestScheduler Policy¶

2.4 Engine-Driven Execution Loop¶

3. Worker¶

Architecture¶

Key Components¶

3.1 Worker Process Structure¶

3.2 GPU Worker¶

3.3 Worker Busy Loop¶

3.4 Model Execution¶

4. Diffusion Pipeline¶

5. Acceleration Components¶

5.1 Attention Backends¶

Architecture¶

Backend Selection¶

Backend Selection Mechanism¶

Attention Backend Registry¶

Attention Backend Integration¶

5.2 Parallel Attention¶

Architecture¶

Ulysses Sequence Parallelism (USP)¶

Ring Sequence Parallelism¶

Relationship with AttentionBackend¶

5.3 Cache Backends¶

Architecture¶

Cache Backend Interface¶

Available Backends¶

Cache Backend Selector¶

5.4 Parallel Strategies¶

Parallelism Types¶

Parallel Group Management¶

6. Data Flow¶

Complete Request Flow¶

2.2 Shared Bookkeeping in `BaseScheduler`¶

2.3 Current `RequestScheduler` Policy¶