vllm_omni.diffusion.offloader ¶

One ring of repeated blocks that a backend streams together.

attrs names the owner attributes that contributed blocks; backends use it to skip those children when they place the non-streamed remainder. It is filled for DiT stacks only — encoder stacks are hooked as whole stacks and their remainder is placed by tensor identity.

attrs: tuple[str, ...] = ()

blocks: tuple[Module, ...]

resident: tuple[Module, ...]

resident_head: int = 0

streaming: tuple[Module, ...]

Retain reusable allocator blocks without monopolizing device memory.

Component offload normally calls empty_cache after every stage. That makes the next stage return to the device allocator even though PyTorch's cached blocks are immediately reusable. This policy keeps the cache while both of these bounds hold:

• cached-but-unallocated memory is at most 25% of device capacity; and
• at least 5% of device capacity is physically free.

Missing memory telemetry is handled conservatively by releasing the cache. Failure paths can force release; normal executor shutdown keeps its own unconditional device-cache cleanup because this policy is not global.

device = device

max_cached_fraction = max_cached_fraction

min_free_fraction = min_free_fraction

release_if_needed(*, force: bool = False) -> bool

Bases: OffloadBackend

Distributed layer-wise (block-level) offloading backend.

Supports both GPU (CUDA) and NPU (CANN) platforms. Device type is determined by the device passed to enable().

Each rank stores only a shard of each block's weights on host memory. Two device slots alternate: one holds current weights, one holds next weights. H2D and AllGather run asynchronously on dedicated streams, overlapped with computation.

comm_stream = current_omni_platform.Stream()

copy_stream = current_omni_platform.Stream()

dp_group: ProcessGroup | None = None

dp_size = config.dp_size

host_weight_plan = host_weight_plan

rank = 0

disable() -> None

enable(pipeline: Module) -> None

load_resident_layers() -> None

offload_resident_layers() -> None

shutdown() -> None

Bases: ModelHook

Hook for distributed layerwise offloading with fixed double-buffer.

Each rank stores only a shard of each block's weights on host memory. Two device slots alternate: one holds current weights, one holds next weights. H2D and AllGather run asynchronously on dedicated streams, overlapped with computation.

Supports both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms.

block_buffers: dict[str, Tensor] = {}

block_parameters: dict[str, Parameter] = {}

comm_stream = comm_stream or current_omni_platform.Stream()

copy_stream = copy_stream or current_omni_platform.Stream()

cpu_shards: dict[dtype, Tensor] = {}

cpu_sources: dict[dtype, list[dict[str, Any]]] = {}

cpu_staging_buffers: list[dict[dtype, Tensor] | None] = [
    None,
    None,
]

cpu_staging_events: list[Any | None] = [None, None]

current_slot = 0

device = device

dp_group = dp_group

dp_size = dp_size

gpu_buffers: list[dict[dtype, Tensor] | None] = (
    shared_buffers
)

gpu_shard_buffers: list[dict[dtype, Tensor] | None] = [
    None,
    None,
]

is_materialized: bool

metadata: dict[dtype, list[dict[str, Any]]] = {}

next_block = next_block

next_block_buffers: dict[str, Tensor] = {}

next_block_parameters: dict[str, Parameter] = {}

pin_memory = pin_memory

rank = rank

rank_local_mmap = rank_local_mmap

ready_events: list[Any | None] = [None, None]

registered_mmap = False

tensor_transforms = tensor_transforms or {}

initialize_hook(module: Module) -> Module

offload_layer() -> None

post_forward(module: Module, output: Any) -> Any

pre_forward(
    module: Module, *args: Any, **kwargs: Any
) -> tuple[tuple, dict]

prefetch_layer(
    slot: int, non_blocking: bool = True
) -> None

restore_next_block_to_cpu() -> None

Bases: OffloadBackend

Layer-wise (block-level) offloading backend.

Implements sliding window offloading where only a small number of transformer blocks reside on GPU at a time. Blocks are prefetched asynchronously while previous blocks compute, and freed after use.

copy_stream = current_omni_platform.Stream()

disable() -> None

enable(pipeline: Module) -> None

Bases: OffloadBackend

Model-level (sequential) offloading backend.

Uses SequentialOffloadHook registered via HookRegistry for automatic module swapping.

disable() -> None

enable(pipeline: Module) -> None

Bases: ABC

Base class for CPU offload backends

config = config

device = device

enabled = False

disable() -> None

enable(pipeline: Module) -> None

is_enabled() -> bool

shutdown() -> None

Optional declarative metadata for component layerwise offload.

Models declare this as a class attribute _offload_plan on the pipeline class. When present, both layerwise backends use it instead of model-specific backend branches, making new integrations data-driven.

If not declared, the offloader falls back to _layerwise_offload_blocks_attrs on each DiT module class.

:func:~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan is the only consumer; backends read the resolved artifact it returns.

Attributes:

block_attrs: dict[str, tuple[str, ...]] = field(
    default_factory=dict
)

encoder_block_attrs: dict[str, tuple[str, ...]] = field(
    default_factory=dict
)

encoder_component_types: dict[str, str] = field(
    default_factory=dict
)

encoder_dlo_weight_replication: frozenset[str] = field(
    default_factory=frozenset
)

offload_submodules: dict[str, str] = field(
    default_factory=dict
)

on_demand_component_paths: frozenset[str] = field(
    default_factory=frozenset
)

resident_dit_paths: frozenset[str] = field(
    default_factory=frozenset
)

Loader-owned state consumed by the offloader startup boundary.

allow_fresh_retry: bool = False

fresh_model_loader: Callable[[], Module] | None = None

host_weight_plan: HostWeightPlan | None = None

close_loader_ownership() -> None

Bases: str, Enum

Resolved internal backend strategy.

DISTRIBUTED_LAYER_WISE = 'distributed_layer_wise'

LAYER_WISE = 'layer_wise'

MODEL_LEVEL = 'model_level'

NONE = 'none'

Stage immutable module groups without copying device weights to CPU.

nn.Module.to("cpu") performs a device-to-host copy for every parameter after each forward. In inference the weights are immutable, so retain one pinned CPU master instead. load materializes device storage from that master; offload only rebinds Parameters and buffers to the master.

A module iterable is treated as one staging group. It uses one copy stream and one reusable completion event. Tensors sharing storage keep their shapes, strides, offsets, dtypes, and aliases across every transition.

cache_retention = cache_retention

copy_stream = (
    copy_stream
    if copy_stream is not None
    else current_omni_platform.Stream()
)

device = device

loaded = False

load() -> None

offload() -> None

set_cache_retention(
    cache_retention: BoundedAllocatorCache | None,
) -> None

One pipeline component with its resolved, validated offload topology.

selected means the active selector covers this component. on_demand means the pipeline owns its residency through load_to_device / offload_to_cpu; a VAE staged by a legacy model plan is on_demand without being selectable through the public component grammar.

children: tuple[ResolvedComponent, ...] = ()

module: Module

on_demand: bool = False

path: str

selected: bool

stacks: tuple[BlockStack, ...] = ()

Backend-neutral topology for one pipeline under one offload config.

components: tuple[ResolvedComponent, ...]

dits: tuple[ResolvedComponent, ...]

encoders: tuple[ResolvedComponent, ...]

residents: tuple[ResolvedComponent, ...]

skip_reason: str | None = None

vaes: tuple[ResolvedComponent, ...]

Bases: Protocol

Pipeline-owned lifecycle for model-level CPU offload.

Pipelines with non-forward component entry points (for example VAE decode_latent methods) need to activate those stages explicitly, so generic forward-hook discovery cannot manage their full lifecycle.

disable_omni_model_cpu_offload() -> None

enable_omni_model_cpu_offload(
    *,
    device: device,
    pin_memory: bool,
    use_hsdp: bool,
    offload_components: frozenset[str] | None = None,
) -> None