Skip to content

vllm_omni.diffusion.offloader.offload_plan

Declarative model capabilities for component layerwise offload.

Models declare an _offload_plan class attribute on the pipeline. Both the ordinary and distributed layerwise backends consume the same plan instead of embedding model-specific block paths in backend code.

OffloadPlan dataclass

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:

Name Type Description
block_attrs dict[str, tuple[str, ...]]

Maps DiT path → tuple of block-list attribute names. e.g. {"transformer": ("gen_layers",), "transformer.language_model": ("layers",)}
offload_submodules dict[str, str]

Maps child name → block-list attribute name, for large non-DiT submodules within a DiT that should be independently offloaded with their own hooks. e.g. {"context_encoder": "layers"}
resident_dit_paths frozenset[str]

DiT paths whose leading blocks may be kept on the device when dlo_resident_layers is nonzero. Keeping this model-declared avoids applying a consumer-GPU tuning knob to auxiliary or dual DiTs unintentionally.
encoder_component_types dict[str, str]

Maps encoder paths to public selector types (currently text_encoder). This declaration is used before the compatibility name heuristic.
encoder_block_attrs dict[str, tuple[str, ...]]

Maps encoder paths to streamable block-list paths.
encoder_dlo_weight_replication frozenset[str]

Encoder paths whose loader-produced block tensors are identical across the DiT DLO group. Only these encoders may use multi-rank AllGather transfer; this must not be declared for encoder-TP shards.

block_attrs class-attribute instance-attribute 

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

encoder_block_attrs class-attribute instance-attribute 

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

encoder_component_types class-attribute instance-attribute 

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

encoder_dlo_weight_replication class-attribute instance-attribute 

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

offload_submodules class-attribute instance-attribute 

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

on_demand_component_paths class-attribute instance-attribute 

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

resident_dit_paths class-attribute instance-attribute 

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

get_offload_plan 

get_offload_plan(pipeline: Module) -> OffloadPlan | None

Retrieve the OffloadPlan declared by the pipeline, if any.