Skip to content

vllm_omni.diffusion.model_loader.host_weights

Typed diffusion adapters for runtime-ready final-layout artifacts.

Modules:

Name Description
contracts

Typed semantic contracts shared by final-layout artifact representations.
identity_adapter

Exact pre-load identity assembly for policy-defined final-layout artifacts.
producers

Diffusion Host Weight Runtime producers.
restorer

Generic validation-first restorer for exact final-layout tensor artifacts.
source_identity

Canonical diffusion weight-source identity and one-load replacement guard.
tensor_layout

Representation-neutral ownership discovery for final-layout tensors.

FINAL_LAYOUT_BF16_POLICY module-attribute 

FINAL_LAYOUT_BF16_POLICY = FinalLayoutBF16Policy()

FINAL_LAYOUT_BF16_SPEC module-attribute 

FINAL_LAYOUT_BF16_SPEC = FinalLayoutArtifactSpec(
    representation=WeightRepresentation(
        name=FINAL_LAYOUT_BF16_REPRESENTATION,
        dtype=str(torch.bfloat16),
        metadata=CanonicalJson.from_value(
            {
                "format": FINAL_LAYOUT_BF16_REPRESENTATION,
                "mixed_precision_policy": "bf16-with-preserved-fp32",
            }
        ),
    ),
    producer=ProducerIdentity(
        producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID,
        version=FINAL_LAYOUT_BF16_VERSION,
        implementation_fingerprint=implementation_abi_fingerprint(
            _BF16_IMPLEMENTATION_ABI
        ),
        manifest_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA,
        restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
    ),
    implementation_abi=_BF16_IMPLEMENTATION_ABI,
)

FinalLayoutArtifactSpec dataclass

Concrete representation, layout, producer, and restoration ABI.

component_name class-attribute instance-attribute 

component_name: str = 'diffusion-dit'

component_ownership class-attribute instance-attribute 

component_ownership: str = 'complete-final-layout-tensors'

implementation_abi instance-attribute 

implementation_abi: CanonicalJson

layout_name class-attribute instance-attribute 

layout_name: str = 'diffusion-final-module-layout-v1'

model_contract_schema class-attribute instance-attribute 

model_contract_schema: str = (
    FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
)

producer instance-attribute 

producer: ProducerIdentity

representation instance-attribute 

representation: WeightRepresentation

FinalLayoutBF16Producer

Publish finalized BF16-policy tensors through a store-scoped writer.

spec property 

spec: WeightProductionSpec

produce 

produce(writer: ArtifactWriter) -> ProductionMetadata

FinalLayoutIdentityContext dataclass

Exact semantic identity plus process-local transaction guards.

dit_names instance-attribute 

dit_names: tuple[str, ...]

dit_types instance-attribute 

dit_types: tuple[type[Module], ...]

identity instance-attribute 

identity: WeightArtifactIdentity

model_contracts instance-attribute 

model_contracts: tuple[FinalLayoutModelContract, ...]

pipeline_type instance-attribute 

pipeline_type: type[Module]

policy instance-attribute 

policy: FinalLayoutTensorPolicy

request instance-attribute 

request: FinalLayoutRequest

source instance-attribute 

source: FinalLayoutSourceContext

spec instance-attribute 

spec: FinalLayoutArtifactSpec

tensor_names instance-attribute 

tensor_names: frozenset[str]

ensure_sources_unchanged 

ensure_sources_unchanged() -> None

sources_unchanged 

sources_unchanged() -> bool

FinalLayoutLoaderIdentity dataclass

Loader ABI and exact configuration facts that may change final bytes.

implementation instance-attribute 

implementation: ImplementationIdentity

model_config_fingerprint instance-attribute 

model_config_fingerprint: str

weight_transform_fingerprint instance-attribute 

weight_transform_fingerprint: str

to_dict 

to_dict() -> dict[str, object]

FinalLayoutParallelIdentity dataclass

Semantic parallel layout with DP and SP rank intentionally absent.

allgather_degree class-attribute instance-attribute 

allgather_degree: int = 1

cfg_parallel_size class-attribute instance-attribute 

cfg_parallel_size: int = 1

enable_expert_parallel class-attribute instance-attribute 

enable_expert_parallel: bool = False

pipeline_parallel_size class-attribute instance-attribute 

pipeline_parallel_size: int = 1

ring_degree class-attribute instance-attribute 

ring_degree: int = 1

sequence_parallel_backend property 

sequence_parallel_backend: str

sequence_parallel_size class-attribute instance-attribute 

sequence_parallel_size: int = 1

tensor_parallel_rank class-attribute instance-attribute 

tensor_parallel_rank: int = 0

tensor_parallel_size class-attribute instance-attribute 

tensor_parallel_size: int = 1

ulysses_degree class-attribute instance-attribute 

ulysses_degree: int = 1

ulysses_mode class-attribute instance-attribute 

ulysses_mode: str = 'strict'

use_hsdp class-attribute instance-attribute 

use_hsdp: bool = False

to_dict 

to_dict() -> dict[str, object]

FinalLayoutRequest dataclass

Representation-independent semantic request assembled by the loader.

adaptation class-attribute instance-attribute 

adaptation: AdaptationIdentity = field(
    default_factory=AdaptationIdentity
)

load_format class-attribute instance-attribute 

load_format: str = 'default'

loader instance-attribute 

loader: FinalLayoutLoaderIdentity

model_id instance-attribute 

model_id: str

parallel class-attribute instance-attribute 

parallel: FinalLayoutParallelIdentity = field(
    default_factory=FinalLayoutParallelIdentity
)

FinalLayoutTensorRestorer

Validate and plan restoration for one exact policy-defined lease.

schema property 

schema: str

plan_restore 

plan_restore(
    model: object, lease: HostWeightLease
) -> FinalLayoutTensorRestorePlan

ImplementationIdentity dataclass

Explicit stable identity for one semantic implementation boundary.

fingerprint instance-attribute 

fingerprint: str

implementation_id instance-attribute 

implementation_id: str

version instance-attribute 

version: str

to_dict 

to_dict() -> dict[str, str]

NodeSourceDigestCache

Share content-backed local-source fingerprints across node workers.

entries instance-attribute 

entries = root / 'source-digests-v1' / 'entries'

locks instance-attribute 

locks = root / 'source-digests-v1' / 'locks'

timeout_seconds instance-attribute 

timeout_seconds = timeout_seconds

content_id 

content_id(path: Path, state: _FileState) -> str

PreparedWeightSource dataclass

Canonical source files and adapter ABI selected by the ordinary loader.

checkpoint_adapter class-attribute instance-attribute 

checkpoint_adapter: ImplementationIdentity | None = None

model_or_path instance-attribute 

model_or_path: str

prefix instance-attribute 

prefix: str

requested_revision instance-attribute 

requested_revision: str | None

resolved_root instance-attribute 

resolved_root: Path

source_kind class-attribute instance-attribute 

source_kind: WeightSourceKind = WeightSourceKind.LOCAL_PATH

subfolder instance-attribute 

subfolder: str | None

use_safetensors instance-attribute 

use_safetensors: bool

weight_files instance-attribute 

weight_files: tuple[Path, ...]

WeightSourceKind

Bases: str, Enum

HUGGING_FACE_HUB class-attribute instance-attribute 

HUGGING_FACE_HUB = 'hugging_face_hub'

LOCAL_PATH class-attribute instance-attribute 

LOCAL_PATH = 'local_path'

build_final_layout_identity 

build_final_layout_identity(
    pipeline: Module,
    *,
    dit_modules: Sequence[tuple[str, Module]],
    prepared_sources: Sequence[PreparedWeightSource],
    request: FinalLayoutRequest,
    policy: FinalLayoutTensorPolicy,
    source_digest_cache: NodeSourceDigestCache
    | None = None,
) -> FinalLayoutIdentityContext

Build exact source, representation, layout, and implementation identity.