vllm_omni.diffusion.offloader.layerwise_backend ¶

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 = Stream()

disable() -> None

enable(pipeline: Module) -> None

get_blocks_attr_names(model: Module) -> list[str]

get_blocks_from_dit(
    model: Module,
) -> tuple[list[str], list[Module]]

class WanTransformer3DModel(nn.Module):
    _layerwise_offload_blocks_attrs = ["blocks"]

set_blocks_attr_names(
    model: Module, names: list[str]
) -> None

Bases: ModelHook

Hook for layerwise (transformer-block-wise) CPU offloading.

The hook instance retains parameters for both the current registered block module and those for the next block, as well as flattened CPU tensors which record the parameters of the current block module, so that these parameters could be re-materialized on device in an overlapping way. This hook should be registered to each of the transformer blocks in DiT module(s) of the target pipeline.

Based on implementations from: https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py

copy_stream = stream or current_stream()

device = device

dtype_cpu_flattened_weights: dict[dtype, Tensor] = {}

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

is_materialized: bool

next_block = next_block

next_block_buffers: dict[str, Tensor] = {}

next_block_parameters: dict[str, Parameter] = {}

pin_memory = pin_memory

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(non_blocking: bool = True) -> None