vllm_omni.diffusion.offloader.block_discovery ¶

Block discovery for layerwise offload.

Shared between LayerWiseOffloadBackend and DistributedLayerwiseOffloadBackend.

logger `module-attribute` ¶

logger = init_logger(__name__)

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

Get block attribute names from model class.

get_blocks_from_dit(
    model: Module,
    block_attrs: tuple[str, ...] | None = None,
) -> tuple[list[str], list[Module]]

Retrieve blocks from an explicit plan or the DiT class metadata.

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