vllm-omni/docs/design/module/archive/dit_module.md

---
toc_depth: 4
---

# 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.

<p align="center">
   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
</p>
<p align="center">
  <em> Main Components of the Diffusion Module </em>
</p>


**Table of Content:**

- [Architecture Overview](#architecture-overview)
- [Diffusion Engine](#1-diffusion-engine)
- [Scheduler](#2-scheduler)
- [Worker](#3-worker)
- [Diffusion Pipeline](#4-diffusion-pipeline)
- [Acceleration Components](#5-acceleration-components)
    - [Attention Backends](#51-attention-backends)
    - [Parallel Attention](#52-parallel-attention)
    - [Cache Backends](#53-cache-backends)
    - [Parallel Strategies](#54-parallel-strategies)
- [Data Flow](#6-data-flow)

---

## Architecture Overview

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

<p align="center">
  <img src="https://github.com/user-attachments/assets/78c4d446-d238-4406-a057-eb17d2ccc8e0" alt="vLLM-Omni Diffusion Module Architecture" width="100%">
</p>
<p align="center">
  <em> Diffusion Architecture Overview </em>
</p>


---

## 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

```python
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

```python
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

```python
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

```python
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**:

1. **Pre-processing**: Applies model-specific transformations

2. **Scheduling**: Delegates request state transitions to the selected scheduler

3. **Execution**: Calls the selected executor function for request, request-batch, or step execution

4. **Output Delivery**: Emits intermediate chunks and final outputs through the request output stream

5. **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

```python
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`

```python
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

```python
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

```python
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](../../feature/diffusion_continuous_batching.md).

---

## 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

```python
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**:

1. **IPC Setup**: Creates ZMQ context and message queues

2. **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.)

3. **Model Loading**: Loads diffusion pipeline on assigned GPU

4. **Cache Setup**: Enables cache backend if configured.

#### 3.2 GPU Worker

```python
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

```python
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**:

1. **Receive**: Dequeues unified messages from shared memory queue

2. **Route**: Handles different message types (generation, RPC, shutdown)

3. **Execute**: Runs forward pass through pipeline for generation requests

4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)

#### 3.4 Model Execution

```python
@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:

```python
# 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:

```python
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](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/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`

```python
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](#52-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:

```python
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**:

1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.

2. **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"`).

3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).

4. **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](../../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../../contributing/model/adding_diffusion_model.md).

**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)`.

```python
@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:

```python
# 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:

1. **Sequence Dimension**: Splits the sequence length across GPUs

2. **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**:
```python
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](#54-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

```python
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`

```python
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/`

```python
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`

```python
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**:

1. **Selection**: `get_cache_backend()` returns appropriate backend instance

2. **Enable**: `backend.enable(pipeline)` called during worker initialization

3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation

4. **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](#52-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

```python
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](#52-parallel-attention). This section covers higher-level parallelism strategies.


---

## 6. Data Flow

### Complete Request Flow

<p align="center">
   <img src="https://github.com/user-attachments/assets/6e093a0d-29c0-4efd-85d5-747002bd2fed" alt="vLLM-Omni Diffusion Module Components" width="100%">
</p>
<p align="center">
  <em> End-to-end Data Flow in the vLLM-Omni Diffusion Module </em>
</p>


```
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
```

---