Skip to content

vllm_omni.diffusion.model_loader.host_weight_plan

Loader-owned host-weight plans shared with diffusion offload backends.

TensorTransform module-attribute 

TensorTransform = Callable[[torch.Tensor], torch.Tensor]

logger module-attribute 

logger = init_logger(__name__)

HostWeightPlan dataclass

Complete, prevalidated host backing consumed by an offload backend.

backing_kind instance-attribute 

backing_kind: str

bindings instance-attribute 

bindings: dict[str, TensorBinding]

lease_carrier class-attribute instance-attribute 

lease_carrier: HostWeightLeaseCarrier | None = None

planned_source_prefixes class-attribute instance-attribute 

planned_source_prefixes: frozenset[str] = frozenset()

HostWeightPlanResult dataclass

A complete plan or a fail-closed reason to use the ordinary loader.

fallback_reason class-attribute instance-attribute 

fallback_reason: str | None = None

plan instance-attribute 

plan: HostWeightPlan | None

TensorBinding dataclass

One runtime tensor backed by one safetensors entry.

checkpoint_key instance-attribute 

checkpoint_key: str

file_path instance-attribute 

file_path: str

transform class-attribute instance-attribute 

transform: TensorTransform | None = None

build_checkpoint_mmap_plan 

build_checkpoint_mmap_plan(
    pipeline: Module,
    *,
    dit_modules: Sequence[tuple[str, Module]],
    sources: Sequence[object],
    model_path: str | None,
    tensor_parallel_size: int,
    use_hsdp: bool,
    online_quantization: bool,
    has_distilled_lora: bool = False,
) -> HostWeightPlanResult

Build a complete direct-checkpoint plan or return a fallback reason.

has_online_quantization 

has_online_quantization(model: Module) -> bool

Whether any module defers online-quantized weights on meta.