[Diffusion] Add final-layout BF16 Host Weight Runtime artifacts (#6445)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/host_weight_runtime.md

@@ -15,8 +15,11 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
 ## Status
 
 The first implementation provides contracts and a CPU local-filesystem store.
-It does not enable cached loading for a model by itself. A consumer must still
-provide an exact identity, a representation producer, and a model restorer.
+The initial diffusion consumer contract additionally defines typed,
+representation-independent final-layout identity/restoration mechanics plus a
+concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
+library-only: no loader or DLO path selects, publishes, restores, or transports
+that artifact yet.
 
 V1 includes:
 
@@ -30,8 +33,9 @@ V1 includes:
 
 V1 does not include:
 
-- a public CLI or default-on model integration;
-- a generic BF16, FP8, quantized, or model-specific producer;
+- a public CLI, loader activation, or default-on model integration;
+- lease handoff to DLO or another transport consumer;
+- FP8, quantized, merged-adaptation, or additional model producers;
 - CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
 - a remote artifact provider or cross-node coordination;
 - automatic eviction; or
@@ -216,6 +220,56 @@ Dynamic LoRA overlays are not part of a reusable base-weight artifact. A
 statically merged adapter is cacheable only as a separate identity containing
 the adapter fingerprint and merge semantics.
 
+### Initial diffusion final-layout contract
+
+The shared diffusion contract covers complete final-layout DiT parameters and
+persistent buffers. Text encoders, VAEs, non-persistent derived state, and other
+pipeline components remain outside the artifact. One explicit representation
+policy selects allowed dtypes, tensor roles, physical layout identity, producer
+ABI, manifest schema, and restoration schema.
+
+The contract is intentionally separate from loader activation:
+
+- `FinalLayoutRequest` contains typed loader identity/configuration fingerprints,
+  TP coordinate, and conservative SP semantics. It has no open metadata bag,
+  DP coordinate, SP rank, device identity, DLO transfer mode, registration
+  policy, or store path.
+- `FinalLayoutArtifactSpec` binds one `WeightRepresentation` and runtime-layout
+  name to explicit producer/restorer schemas and a canonical, versioned
+  implementation ABI descriptor. Compatibility never depends on reflective
+  source inspection.
+- `PreparedWeightSource` snapshots immutable revisions or exact local file
+  content plus a typed checkpoint-adapter identity before ordinary
+  materialization. Source replacement before or during production fails
+  publication. A hash-looking symlink basename is trusted only for an explicit
+  Hugging Face Hub source whose repository ID and
+  `models--.../snapshots/<revision> -> blobs/<hash>` topology validate; every
+  local or otherwise unverified symlink target is content-hashed.
+- the tensor ownership digest records exact runtime names, kinds, shapes,
+  semantic roles, dtypes, and strides from a CPU or meta model skeleton;
+- `FinalLayoutTensorRestorer` accepts only an exact lease identity, validates
+  complete policy-defined coverage without mutation, and returns a one-shot
+  commit plan;
+- each model declares one dtype-neutral `FinalLayoutModelContract` with an
+  explicit implementation version and a post-commit validator; and
+- `FinalLayoutBF16Producer` accepts only the matching identity context and a
+  finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
+  coordinate. Its BF16 policy preserves model-declared FP32 parameters and
+  buffers and revalidates MiniMax H3 mixed-precision invariants.
+
+Other representations reuse source identity, typed parallel identity, tensor
+ownership, and exact restoration only when their policy proves those semantics.
+For example, runtime FP8 needs a separate policy/producer for generated scales,
+quantization metadata, and Cutlass physical layouts; it is not enabled by
+changing a dtype string on the BF16 producer.
+
+This stage makes no startup, sharing, or DLO performance claim. A following
+consumer PR owns disabled/preferred/required precedence, mixed-component loader
+transactions, warm-hit restoration, and transactional lease handoff. A TP2
+prewarm deployment will require a matching TP2 producer cohort to populate both
+TP-coordinate identities even though the store coordinates each artifact
+independently.
+
 ## Host sharing and GPU transport
 
 Every process receives its own virtual mappings and `HostWeightLease`, but

docs/design/module/host_weight_runtime.md

@@ -56,7 +56,8 @@ The contracts have deliberately narrow ownership:
 
 - `HostWeightRuntime` applies resolution and fallback policy and emits one
   terminal resolution report. Explicit post-load publication emits a separate
-  report and may return a lease, but does not revise the completed resolution.
+  report and closes its store-returned lease before returning; it does not
+  revise the completed resolution.
   The runtime does not own a model or call the canonical loader.
 - `HostWeightStore` owns artifact lookup, coordinated construction,
   publication, validation, quarantine, and lifecycle.

vllm_omni/diffusion/model_loader/host_weights/__init__.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Typed diffusion adapters for runtime-ready final-layout artifacts."""
+
+from .contracts import (
+    FinalLayoutArtifactSpec,
+    FinalLayoutLoaderIdentity,
+    FinalLayoutParallelIdentity,
+    FinalLayoutRequest,
+    ImplementationIdentity,
+)
+from .identity_adapter import (
+    FinalLayoutIdentityContext,
+    build_final_layout_identity,
+)
+from .producers import (
+    FINAL_LAYOUT_BF16_POLICY,
+    FINAL_LAYOUT_BF16_SPEC,
+    FinalLayoutBF16Producer,
+)
+from .restorer import FinalLayoutTensorRestorer
+from .source_identity import PreparedWeightSource, WeightSourceKind
+
+__all__ = [
+    "FINAL_LAYOUT_BF16_POLICY",
+    "FINAL_LAYOUT_BF16_SPEC",
+    "FinalLayoutArtifactSpec",
+    "FinalLayoutBF16Producer",
+    "FinalLayoutIdentityContext",
+    "FinalLayoutLoaderIdentity",
+    "FinalLayoutParallelIdentity",
+    "FinalLayoutRequest",
+    "FinalLayoutTensorRestorer",
+    "ImplementationIdentity",
+    "PreparedWeightSource",
+    "WeightSourceKind",
+    "build_final_layout_identity",
+]

vllm_omni/diffusion/model_loader/host_weights/contracts.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Typed semantic contracts shared by final-layout artifact representations."""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Protocol
+
+import torch
+
+from vllm_omni.diffusion.models.host_weight_contract import (
+    FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+)
+from vllm_omni.host_weight_runtime import (
+    AdaptationIdentity,
+    CanonicalJson,
+    FailureCode,
+    HostWeightError,
+    HostWeightFailure,
+    ProducerIdentity,
+    ResolutionStage,
+    TensorKind,
+    WeightRepresentation,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from .tensor_layout import RuntimeTensorTarget
+
+FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA = "diffusion-final-layout-tensor-restorer-v1"
+
+
+def _require_text(name: str, value: object) -> str:
+    if not isinstance(value, str) or not value:
+        raise ValueError(f"{name} must not be empty")
+    return value
+
+
+def _require_positive_int(name: str, value: object) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValueError(f"{name} must be a positive integer")
+    return value
+
+
+class FinalLayoutContractCode(str, Enum):
+    SOURCE_CHANGED = "source_changed"
+    SOURCE_COVERAGE_INVALID = "source_coverage_invalid"
+    MODEL_CONTRACT_UNSUPPORTED = "model_contract_unsupported"
+    OWNERSHIP_AMBIGUOUS = "ownership_ambiguous"
+    TENSOR_UNSUPPORTED = "tensor_unsupported"
+    DTYPE_UNSUPPORTED = "dtype_unsupported"
+    IDENTITY_INCOMPATIBLE = "identity_incompatible"
+    TENSOR_CONTRACT_CHANGED = "tensor_contract_changed"
+
+
+class FinalLayoutContractError(ValueError):
+    """Typed fail-closed error at a final-layout model integration boundary."""
+
+    def __init__(self, code: FinalLayoutContractCode, message: str) -> None:
+        if not isinstance(code, FinalLayoutContractCode):
+            raise ValueError("final-layout contract errors require FinalLayoutContractCode")
+        super().__init__(message)
+        self.code = code
+
+
+def final_layout_producer_error(error: FinalLayoutContractError) -> HostWeightError:
+    """Translate representation failures into the runtime's stable taxonomy."""
+    if error.code is FinalLayoutContractCode.SOURCE_CHANGED:
+        stage = ResolutionStage.CANONICAL_LOADING
+        code = FailureCode.CANONICAL_SOURCE_FAILED
+    elif error.code in {
+        FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED,
+        FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+        FinalLayoutContractCode.DTYPE_UNSUPPORTED,
+    }:
+        stage = ResolutionStage.PRODUCTION
+        code = FailureCode.PRODUCER_UNSUPPORTED
+    else:
+        stage = ResolutionStage.PRODUCTION
+        code = FailureCode.PRODUCER_FAILED
+    return HostWeightError(
+        HostWeightFailure(
+            stage=stage,
+            code=code,
+            retryable=False,
+            message=str(error),
+            details=CanonicalJson.from_value({"final_layout_code": error.code.value}),
+        )
+    )
+
+
+@dataclass(frozen=True)
+class ImplementationIdentity:
+    """Explicit stable identity for one semantic implementation boundary."""
+
+    implementation_id: str
+    version: str
+    fingerprint: str
+
+    def __post_init__(self) -> None:
+        _require_text("implementation_id", self.implementation_id)
+        _require_text("implementation version", self.version)
+        _require_text("implementation fingerprint", self.fingerprint)
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "implementation_id": self.implementation_id,
+            "version": self.version,
+            "fingerprint": self.fingerprint,
+        }
+
+
+@dataclass(frozen=True)
+class FinalLayoutLoaderIdentity:
+    """Loader ABI and exact configuration facts that may change final bytes."""
+
+    implementation: ImplementationIdentity
+    model_config_fingerprint: str
+    weight_transform_fingerprint: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.implementation, ImplementationIdentity):
+            raise ValueError("loader implementation must use ImplementationIdentity")
+        _require_text("model config fingerprint", self.model_config_fingerprint)
+        _require_text("weight transform fingerprint", self.weight_transform_fingerprint)
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "implementation": self.implementation.to_dict(),
+            "model_config_fingerprint": self.model_config_fingerprint,
+            "weight_transform_fingerprint": self.weight_transform_fingerprint,
+        }
+
+
+@dataclass(frozen=True)
+class FinalLayoutParallelIdentity:
+    """Semantic parallel layout with DP and SP rank intentionally absent."""
+
+    tensor_parallel_size: int = 1
+    tensor_parallel_rank: int = 0
+    sequence_parallel_size: int = 1
+    ulysses_degree: int = 1
+    ring_degree: int = 1
+    allgather_degree: int = 1
+    ulysses_mode: str = "strict"
+    pipeline_parallel_size: int = 1
+    cfg_parallel_size: int = 1
+    use_hsdp: bool = False
+    enable_expert_parallel: bool = False
+
+    def __post_init__(self) -> None:
+        tp_size = _require_positive_int("tensor_parallel_size", self.tensor_parallel_size)
+        if isinstance(self.tensor_parallel_rank, bool) or not isinstance(self.tensor_parallel_rank, int):
+            raise ValueError("tensor_parallel_rank must be an integer")
+        if not 0 <= self.tensor_parallel_rank < tp_size:
+            raise ValueError("tensor_parallel_rank must be within tensor_parallel_size")
+        for name, value in (
+            ("sequence_parallel_size", self.sequence_parallel_size),
+            ("ulysses_degree", self.ulysses_degree),
+            ("ring_degree", self.ring_degree),
+            ("allgather_degree", self.allgather_degree),
+            ("pipeline_parallel_size", self.pipeline_parallel_size),
+            ("cfg_parallel_size", self.cfg_parallel_size),
+        ):
+            _require_positive_int(name, value)
+        _require_text("ulysses_mode", self.ulysses_mode)
+        if not isinstance(self.use_hsdp, bool) or not isinstance(self.enable_expert_parallel, bool):
+            raise ValueError("HSDP and expert-parallel flags must be booleans")
+
+    @property
+    def sequence_parallel_backend(self) -> str:
+        if self.sequence_parallel_size == 1:
+            return "none"
+        active = [
+            name
+            for name, degree in (
+                ("ulysses", self.ulysses_degree),
+                ("ring", self.ring_degree),
+                ("allgather", self.allgather_degree),
+            )
+            if degree > 1
+        ]
+        return "+".join(active) if active else "replicated-sp"
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "tensor_parallel_size": self.tensor_parallel_size,
+            "tensor_parallel_rank": self.tensor_parallel_rank,
+            "sequence_parallel_size": self.sequence_parallel_size,
+            "sequence_parallel_backend": self.sequence_parallel_backend,
+            "ulysses_degree": self.ulysses_degree,
+            "ring_degree": self.ring_degree,
+            "allgather_degree": self.allgather_degree,
+            "ulysses_mode": self.ulysses_mode,
+            "pipeline_parallel_size": self.pipeline_parallel_size,
+            "cfg_parallel_size": self.cfg_parallel_size,
+            "use_hsdp": self.use_hsdp,
+            "enable_expert_parallel": self.enable_expert_parallel,
+        }
+
+
+@dataclass(frozen=True)
+class FinalLayoutRequest:
+    """Representation-independent semantic request assembled by the loader."""
+
+    model_id: str
+    loader: FinalLayoutLoaderIdentity
+    parallel: FinalLayoutParallelIdentity = field(default_factory=FinalLayoutParallelIdentity)
+    load_format: str = "default"
+    adaptation: AdaptationIdentity = field(default_factory=AdaptationIdentity)
+
+    def __post_init__(self) -> None:
+        _require_text("canonical model ID or path", self.model_id)
+        if not isinstance(self.loader, FinalLayoutLoaderIdentity):
+            raise ValueError("loader must use FinalLayoutLoaderIdentity")
+        if not isinstance(self.parallel, FinalLayoutParallelIdentity):
+            raise ValueError("parallel must use FinalLayoutParallelIdentity")
+        _require_text("load_format", self.load_format)
+        if not isinstance(self.adaptation, AdaptationIdentity):
+            raise ValueError("adaptation must use AdaptationIdentity")
+
+
+def implementation_abi_fingerprint(abi: CanonicalJson) -> str:
+    if not isinstance(abi, CanonicalJson):
+        raise ValueError("implementation ABI must use CanonicalJson")
+    return hashlib.sha256(abi.encoded).hexdigest()
+
+
+@dataclass(frozen=True)
+class FinalLayoutArtifactSpec:
+    """Concrete representation, layout, producer, and restoration ABI."""
+
+    representation: WeightRepresentation
+    producer: ProducerIdentity
+    implementation_abi: CanonicalJson
+    layout_name: str = "diffusion-final-module-layout-v1"
+    component_name: str = "diffusion-dit"
+    component_ownership: str = "complete-final-layout-tensors"
+    model_contract_schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.representation, WeightRepresentation):
+            raise ValueError("artifact representation must use WeightRepresentation")
+        if not isinstance(self.producer, ProducerIdentity):
+            raise ValueError("artifact producer must use ProducerIdentity")
+        if not isinstance(self.implementation_abi, CanonicalJson):
+            raise ValueError("artifact implementation ABI must use CanonicalJson")
+        for name, value in (
+            ("layout_name", self.layout_name),
+            ("component_name", self.component_name),
+            ("component_ownership", self.component_ownership),
+            ("model_contract_schema", self.model_contract_schema),
+        ):
+            _require_text(name, value)
+        abi = self.implementation_abi.to_value()
+        required_abi_fields = {
+            "artifact_identity",
+            "model_contract",
+            "producer",
+            "representation_policy",
+            "restorer",
+            "source_identity",
+            "tensor_contract",
+        }
+        if not isinstance(abi, dict) or not required_abi_fields <= set(abi):
+            raise ValueError(
+                "artifact implementation ABI must version identity, source, tensor, policy, producer, "
+                "restorer, and model contracts"
+            )
+        for name in required_abi_fields:
+            _require_text(f"implementation ABI {name}", abi[name])
+        if self.producer.implementation_fingerprint != implementation_abi_fingerprint(self.implementation_abi):
+            raise ValueError("producer implementation fingerprint differs from its explicit ABI descriptor")
+
+
+class FinalLayoutTensorPolicy(Protocol):
+    """Representation-specific validation over shared final-layout mechanics."""
+
+    @property
+    def spec(self) -> FinalLayoutArtifactSpec: ...
+
+    def validate_request(self, request: FinalLayoutRequest) -> None: ...
+
+    def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: ...
+
+    def validate_target(self, target: RuntimeTensorTarget) -> None: ...
+
+    def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: ...
+
+    def build_format_metadata(
+        self,
+        *,
+        component_names: tuple[str, ...],
+        tensor_contract_digest: str,
+        tensor_count: int,
+    ) -> CanonicalJson: ...
+
+    def validate_format_metadata(
+        self,
+        metadata: CanonicalJson,
+        *,
+        component_names: tuple[str, ...],
+        tensor_contract_digest: str,
+        tensor_count: int,
+    ) -> None: ...
+
+
+__all__ = [
+    "FinalLayoutArtifactSpec",
+    "FinalLayoutContractCode",
+    "FinalLayoutContractError",
+    "FinalLayoutLoaderIdentity",
+    "FinalLayoutParallelIdentity",
+    "FinalLayoutRequest",
+    "FinalLayoutTensorPolicy",
+    "FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA",
+    "ImplementationIdentity",
+    "final_layout_producer_error",
+    "implementation_abi_fingerprint",
+]

vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact pre-load identity assembly for policy-defined final-layout artifacts."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+from torch import nn
+
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
+from vllm_omni.host_weight_runtime import (
+    CanonicalJson,
+    ComponentIdentity,
+    RuntimeWeightLayout,
+    WeightArtifactIdentity,
+    WeightRepresentation,
+)
+
+from .contracts import (
+    FinalLayoutArtifactSpec,
+    FinalLayoutContractCode,
+    FinalLayoutContractError,
+    FinalLayoutRequest,
+    FinalLayoutTensorPolicy,
+)
+from .source_identity import (
+    FinalLayoutSourceContext,
+    PreparedWeightSource,
+    resolve_final_layout_source_identity,
+)
+from .tensor_layout import (
+    RuntimeTensorTarget,
+    collect_final_layout_targets,
+    tensor_contract_sha256,
+    validate_final_layout_model_contract,
+)
+
+
+@dataclass(frozen=True)
+class FinalLayoutIdentityContext:
+    """Exact semantic identity plus process-local transaction guards."""
+
+    identity: WeightArtifactIdentity
+    spec: FinalLayoutArtifactSpec
+    policy: FinalLayoutTensorPolicy
+    request: FinalLayoutRequest
+    tensor_names: frozenset[str]
+    dit_names: tuple[str, ...]
+    model_contracts: tuple[FinalLayoutModelContract, ...]
+    source: FinalLayoutSourceContext
+    pipeline_type: type[nn.Module]
+    dit_types: tuple[type[nn.Module], ...]
+
+    def sources_unchanged(self) -> bool:
+        return self.source.sources_unchanged()
+
+    def ensure_sources_unchanged(self) -> None:
+        self.source.ensure_sources_unchanged()
+
+
+def _require_metadata_object(name: str, metadata: CanonicalJson) -> Mapping[str, object]:
+    value = metadata.to_value()
+    if not isinstance(value, dict):
+        raise ValueError(f"{name} must encode one JSON object")
+    return value
+
+
+def _model_contract_values(contracts: Sequence[FinalLayoutModelContract]) -> list[dict[str, str]]:
+    return [contract.to_dict() for contract in contracts]
+
+
+def build_final_layout_identity(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    prepared_sources: Sequence[PreparedWeightSource],
+    request: FinalLayoutRequest,
+    policy: FinalLayoutTensorPolicy,
+) -> FinalLayoutIdentityContext:
+    """Build exact source, representation, layout, and implementation identity."""
+    if not isinstance(request, FinalLayoutRequest):
+        raise TypeError("final-layout identity requires FinalLayoutRequest")
+    policy.validate_request(request)
+    spec = policy.spec
+    bindings = validate_final_layout_model_contract(
+        dit_modules,
+        expected_schema=spec.model_contract_schema,
+    )
+    structural_targets = collect_final_layout_targets(
+        pipeline,
+        dit_modules,
+        policy=policy,
+        require_materialized=False,
+    )
+    target_names = frozenset(target.name for target in structural_targets)
+    source = resolve_final_layout_source_identity(
+        prepared_sources,
+        model_id=request.model_id,
+        target_names=target_names,
+    )
+    dit_names = tuple(name for name, _ in dit_modules)
+    model_contracts = tuple(binding.contract for binding in bindings)
+
+    representation = WeightRepresentation(
+        name=spec.representation.name,
+        dtype=spec.representation.dtype,
+        metadata=CanonicalJson.from_value(
+            {
+                "artifact_policy": _require_metadata_object(
+                    "artifact representation metadata",
+                    spec.representation.metadata,
+                ),
+                "load_format": request.load_format,
+                "loader": request.loader.to_dict(),
+            }
+        ),
+    )
+    identity = WeightArtifactIdentity(
+        schema_version=1,
+        source=source.identity,
+        component=ComponentIdentity(
+            name=spec.component_name,
+            ownership=spec.component_ownership,
+            metadata=CanonicalJson.from_value(
+                {
+                    "component_names": list(dit_names),
+                    "model_contracts": _model_contract_values(model_contracts),
+                }
+            ),
+        ),
+        representation=representation,
+        layout=RuntimeWeightLayout(
+            name=spec.layout_name,
+            tensor_parallel_size=request.parallel.tensor_parallel_size,
+            tensor_parallel_rank=request.parallel.tensor_parallel_rank,
+            sequence_parallel_size=request.parallel.sequence_parallel_size,
+            sequence_parallel_backend=request.parallel.sequence_parallel_backend,
+            metadata=CanonicalJson.from_value(
+                {
+                    "model_contract_schema": spec.model_contract_schema,
+                    "parallel": request.parallel.to_dict(),
+                    "tensor_contract_sha256": tensor_contract_sha256(structural_targets),
+                }
+            ),
+        ),
+        adaptation=request.adaptation,
+        producer=spec.producer,
+    )
+    return FinalLayoutIdentityContext(
+        identity=identity,
+        spec=spec,
+        policy=policy,
+        request=request,
+        tensor_names=target_names,
+        dit_names=dit_names,
+        model_contracts=model_contracts,
+        source=source,
+        pipeline_type=type(pipeline),
+        dit_types=tuple(type(module) for _, module in dit_modules),
+    )
+
+
+def validate_final_layout_identity(
+    context: FinalLayoutIdentityContext,
+    targets: Sequence[RuntimeTensorTarget],
+) -> str:
+    """Validate a context against current targets and return their digest."""
+    identity = context.identity
+    spec = context.spec
+    request = context.request
+    digest = tensor_contract_sha256(targets)
+    if identity.source != context.source.identity or identity.adaptation != request.adaptation:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity source or adaptation differs from the exact context",
+        )
+    if identity.producer != spec.producer:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity producer differs from the artifact specification",
+        )
+    if identity.representation.name != spec.representation.name or identity.representation.dtype != (
+        spec.representation.dtype
+    ):
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity representation differs from the artifact specification",
+        )
+    if identity.component.name != spec.component_name or identity.component.ownership != spec.component_ownership:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity component ownership differs from the artifact specification",
+        )
+    if identity.layout.name != spec.layout_name:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity runtime layout differs from the artifact specification",
+        )
+
+    component = _require_metadata_object("component metadata", identity.component.metadata)
+    representation = _require_metadata_object("representation metadata", identity.representation.metadata)
+    layout = _require_metadata_object("layout metadata", identity.layout.metadata)
+    if representation != {
+        "artifact_policy": _require_metadata_object(
+            "artifact representation metadata",
+            spec.representation.metadata,
+        ),
+        "load_format": request.load_format,
+        "loader": request.loader.to_dict(),
+    }:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity loader or representation policy differs from the exact context",
+        )
+    if component.get("component_names") != list(context.dit_names) or component.get("model_contracts") != (
+        _model_contract_values(context.model_contracts)
+    ):
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity model ownership differs from the exact context",
+        )
+    if layout.get("model_contract_schema") != spec.model_contract_schema:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity model restore schema differs from the artifact specification",
+        )
+    if (
+        identity.layout.tensor_parallel_size != request.parallel.tensor_parallel_size
+        or identity.layout.tensor_parallel_rank != request.parallel.tensor_parallel_rank
+        or identity.layout.sequence_parallel_size != request.parallel.sequence_parallel_size
+        or identity.layout.sequence_parallel_backend != request.parallel.sequence_parallel_backend
+        or layout.get("parallel") != request.parallel.to_dict()
+    ):
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.IDENTITY_INCOMPATIBLE,
+            "identity parallel semantics differ from the exact context",
+        )
+    if layout.get("tensor_contract_sha256") != digest:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED,
+            "identity tensor contract differs from the current model",
+        )
+    if frozenset(target.name for target in targets) != context.tensor_names:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED,
+            "current tensor ownership differs from the exact identity context",
+        )
+    return digest
+
+
+__all__ = [
+    "FinalLayoutIdentityContext",
+    "build_final_layout_identity",
+    "validate_final_layout_identity",
+]

vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Diffusion Host Weight Runtime producers."""
+
+from .final_layout_bf16 import (
+    FINAL_LAYOUT_BF16_MANIFEST_SCHEMA,
+    FINAL_LAYOUT_BF16_POLICY,
+    FINAL_LAYOUT_BF16_PRODUCER_ID,
+    FINAL_LAYOUT_BF16_REPRESENTATION,
+    FINAL_LAYOUT_BF16_SPEC,
+    FINAL_LAYOUT_BF16_VERSION,
+    FinalLayoutBF16Policy,
+    FinalLayoutBF16Producer,
+)
+
+__all__ = [
+    "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA",
+    "FINAL_LAYOUT_BF16_POLICY",
+    "FINAL_LAYOUT_BF16_PRODUCER_ID",
+    "FINAL_LAYOUT_BF16_REPRESENTATION",
+    "FINAL_LAYOUT_BF16_SPEC",
+    "FINAL_LAYOUT_BF16_VERSION",
+    "FinalLayoutBF16Policy",
+    "FinalLayoutBF16Producer",
+]

vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py

@@ -0,0 +1,315 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Concrete BF16-with-preserved-FP32 final-layout artifact policy and producer."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.host_weight_contract import (
+    FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+)
+from vllm_omni.host_weight_runtime import (
+    ArtifactWriter,
+    CanonicalJson,
+    CoordinationScope,
+    LookupPhase,
+    ProducerIdentity,
+    ProductionMetadata,
+    ProductionSourceMode,
+    TensorKind,
+    TensorWriteSpec,
+    WeightProductionSpec,
+    WeightRepresentation,
+)
+
+from ..contracts import (
+    FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
+    FinalLayoutArtifactSpec,
+    FinalLayoutContractCode,
+    FinalLayoutContractError,
+    FinalLayoutRequest,
+    final_layout_producer_error,
+    implementation_abi_fingerprint,
+)
+from ..identity_adapter import (
+    FinalLayoutIdentityContext,
+    validate_final_layout_identity,
+)
+from ..tensor_layout import (
+    RuntimeTensorTarget,
+    collect_final_layout_targets,
+    validate_final_layout_model_contract,
+)
+
+FINAL_LAYOUT_BF16_PRODUCER_ID = "vllm-omni.diffusion.final-layout-bf16"
+FINAL_LAYOUT_BF16_VERSION = "1"
+FINAL_LAYOUT_BF16_REPRESENTATION = "diffusion-final-layout-bf16"
+FINAL_LAYOUT_BF16_MANIFEST_SCHEMA = "diffusion-final-layout-bf16-manifest-v1"
+DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
+
+_BF16_PARAMETER_DTYPES = {torch.bfloat16, torch.float32}
+_BF16_BUFFER_DTYPES = {
+    torch.bool,
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+    torch.float64,
+}
+for _dtype_name in (
+    "float8_e4m3fn",
+    "float8_e5m2",
+    "float8_e4m3fnuz",
+    "float8_e5m2fnuz",
+):
+    if (_dtype := getattr(torch, _dtype_name, None)) is not None:
+        _BF16_BUFFER_DTYPES.add(_dtype)
+
+_BF16_IMPLEMENTATION_ABI = CanonicalJson.from_value(
+    {
+        "artifact_identity": "diffusion-final-layout-identity-v1",
+        "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+        "producer": "finalized-bf16-tensor-writer-v1",
+        "representation_policy": "bf16-with-preserved-fp32-v1",
+        "restorer": "exact-final-layout-tensor-rebind-v1",
+        "source_identity": "prepared-diffusion-weight-source-v2",
+        "tensor_contract": "complete-strided-tensor-ownership-v1",
+    }
+)
+
+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,
+)
+
+
+class FinalLayoutBF16Policy:
+    """Concrete compatibility and metadata policy for finalized BF16 tensors."""
+
+    @property
+    def spec(self) -> FinalLayoutArtifactSpec:
+        return FINAL_LAYOUT_BF16_SPEC
+
+    def validate_request(self, request: FinalLayoutRequest) -> None:
+        if request.load_format != "default":
+            raise ValueError(f"final-layout BF16 artifacts require load_format='default', got {request.load_format!r}")
+        if request.adaptation.kind != "base" or request.adaptation.fingerprint is not None:
+            raise ValueError("merged or static LoRA requires a representation-specific producer")
+        parallel = request.parallel
+        if parallel.pipeline_parallel_size != 1:
+            raise ValueError("pipeline-parallel ownership is not supported by the final-layout BF16 producer")
+        if parallel.cfg_parallel_size != 1:
+            raise ValueError("CFG-parallel ownership is not supported by the final-layout BF16 producer")
+        if parallel.use_hsdp:
+            raise ValueError("HSDP/DTensor layouts are not supported by the final-layout BF16 producer")
+        if parallel.enable_expert_parallel:
+            raise ValueError("expert-parallel ownership is not supported by the final-layout BF16 producer")
+
+    def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str:
+        del name, tensor
+        if kind is TensorKind.PARAMETER:
+            return "weight"
+        return "persistent_buffer"
+
+    def validate_target(self, target: RuntimeTensorTarget) -> None:
+        if target.kind is TensorKind.PARAMETER and target.tensor.dtype not in _BF16_PARAMETER_DTYPES:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.DTYPE_UNSUPPORTED,
+                f"{target.name!r} must be BF16 or an explicitly preserved FP32 parameter, got {target.tensor.dtype}",
+            )
+        if target.kind is TensorKind.BUFFER and target.tensor.dtype not in _BF16_BUFFER_DTYPES:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.DTYPE_UNSUPPORTED,
+                f"{target.name!r} uses unsupported buffer dtype {target.tensor.dtype}",
+            )
+
+    def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None:
+        if not any(target.kind is TensorKind.PARAMETER and target.tensor.dtype is torch.bfloat16 for target in targets):
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.DTYPE_UNSUPPORTED,
+                "final-layout BF16 representation requires at least one BF16 parameter",
+            )
+
+    def build_format_metadata(
+        self,
+        *,
+        component_names: tuple[str, ...],
+        tensor_contract_digest: str,
+        tensor_count: int,
+    ) -> CanonicalJson:
+        return CanonicalJson.from_value(
+            {
+                "component_names": list(component_names),
+                "format": FINAL_LAYOUT_BF16_REPRESENTATION,
+                "mixed_precision_policy": "bf16-with-preserved-fp32",
+                "tensor_contract_sha256": tensor_contract_digest,
+                "tensor_count": tensor_count,
+                "tensor_layout": "contiguous",
+            }
+        )
+
+    def validate_format_metadata(
+        self,
+        metadata: CanonicalJson,
+        *,
+        component_names: tuple[str, ...],
+        tensor_contract_digest: str,
+        tensor_count: int,
+    ) -> None:
+        expected = self.build_format_metadata(
+            component_names=component_names,
+            tensor_contract_digest=tensor_contract_digest,
+            tensor_count=tensor_count,
+        )
+        if metadata != expected:
+            raise ValueError("lease format metadata differs from the BF16 artifact policy")
+
+
+FINAL_LAYOUT_BF16_POLICY = FinalLayoutBF16Policy()
+
+
+def _split_shards(
+    records: Sequence[RuntimeTensorTarget],
+    max_shard_bytes: int,
+) -> tuple[tuple[RuntimeTensorTarget, ...], ...]:
+    if max_shard_bytes <= 0:
+        raise ValueError("final-layout BF16 shard size must be positive")
+    shards: list[tuple[RuntimeTensorTarget, ...]] = []
+    current: list[RuntimeTensorTarget] = []
+    current_bytes = 0
+    for record in records:
+        if current and current_bytes + record.nbytes > max_shard_bytes:
+            shards.append(tuple(current))
+            current = []
+            current_bytes = 0
+        current.append(record)
+        current_bytes += record.nbytes
+    if current:
+        shards.append(tuple(current))
+    return tuple(shards)
+
+
+class FinalLayoutBF16Producer:
+    """Publish finalized BF16-policy tensors through a store-scoped writer."""
+
+    def __init__(
+        self,
+        context: FinalLayoutIdentityContext,
+        pipeline: nn.Module,
+        dit_modules: Sequence[tuple[str, nn.Module]],
+        *,
+        max_shard_bytes: int = DEFAULT_SHARD_SIZE_BYTES,
+    ) -> None:
+        if max_shard_bytes <= 0:
+            raise ValueError("final-layout BF16 shard size must be positive")
+        if context.spec != FINAL_LAYOUT_BF16_SPEC:
+            raise ValueError("producer requires the exact BF16 artifact specification")
+        self._context = context
+        self._pipeline = pipeline
+        self._dit_modules = tuple(dit_modules)
+        self._max_shard_bytes = max_shard_bytes
+        if tuple(name for name, _ in self._dit_modules) != context.dit_names:
+            raise ValueError("producer DiT components differ from the exact identity context")
+        if type(pipeline) is not context.pipeline_type or tuple(type(module) for _, module in self._dit_modules) != (
+            context.dit_types
+        ):
+            raise ValueError("producer model implementation differs from the exact identity context")
+        self._spec = WeightProductionSpec(
+            producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID,
+            outputs=(context.identity,),
+            source_mode=ProductionSourceMode.FINALIZED_MODEL,
+            coordination_scope=CoordinationScope.SINGLE_PROCESS,
+            lookup_phase=LookupPhase.POST_LOAD_ONLY,
+        )
+
+    @property
+    def spec(self) -> WeightProductionSpec:
+        return self._spec
+
+    def produce(self, writer: ArtifactWriter) -> ProductionMetadata:
+        try:
+            return self._produce(writer)
+        except FinalLayoutContractError as exc:
+            raise final_layout_producer_error(exc) from exc
+
+    def _produce(self, writer: ArtifactWriter) -> ProductionMetadata:
+        self._context.ensure_sources_unchanged()
+        bindings = validate_final_layout_model_contract(
+            self._dit_modules,
+            expected_schema=self._context.spec.model_contract_schema,
+        )
+        records = collect_final_layout_targets(
+            self._pipeline,
+            self._dit_modules,
+            policy=FINAL_LAYOUT_BF16_POLICY,
+            require_materialized=True,
+        )
+        contract_digest = validate_final_layout_identity(self._context, records)
+        for binding in bindings:
+            binding.validator()
+
+        shards = _split_shards(records, self._max_shard_bytes)
+        shard_count = len(shards)
+        for index, shard in enumerate(shards, start=1):
+            specs = tuple(
+                TensorWriteSpec(
+                    name=record.name,
+                    shape=tuple(record.tensor.shape),
+                    dtype=record.tensor.dtype,
+                    kind=record.kind,
+                    role=record.role,
+                )
+                for record in shard
+            )
+            file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
+            with writer.open_tensor_file(file_name, specs) as output:
+                for record in shard:
+                    output.write_tensor(record.name, record.tensor.detach())
+
+        self._context.ensure_sources_unchanged()
+        return ProductionMetadata(
+            producer_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA,
+            restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
+            format_metadata=FINAL_LAYOUT_BF16_POLICY.build_format_metadata(
+                component_names=self._context.dit_names,
+                tensor_contract_digest=contract_digest,
+                tensor_count=len(records),
+            ),
+        )
+
+
+__all__ = [
+    "DEFAULT_SHARD_SIZE_BYTES",
+    "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA",
+    "FINAL_LAYOUT_BF16_POLICY",
+    "FINAL_LAYOUT_BF16_PRODUCER_ID",
+    "FINAL_LAYOUT_BF16_REPRESENTATION",
+    "FINAL_LAYOUT_BF16_SPEC",
+    "FINAL_LAYOUT_BF16_VERSION",
+    "FinalLayoutBF16Policy",
+    "FinalLayoutBF16Producer",
+]

vllm_omni/diffusion/model_loader/host_weights/restorer.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Generic validation-first restorer for exact final-layout tensor artifacts."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from vllm_omni.host_weight_runtime import HostWeightLease, TensorKind, WeightRestorePlan
+
+from .identity_adapter import (
+    FinalLayoutIdentityContext,
+    validate_final_layout_identity,
+)
+from .tensor_layout import collect_final_layout_targets, validate_final_layout_model_contract
+
+
+@dataclass(frozen=True)
+class _Replacement:
+    name: str
+    parent: nn.Module
+    leaf_name: str
+    target: torch.Tensor
+    source: torch.Tensor
+    kind: TensorKind
+
+    def current_target(self) -> torch.Tensor | None:
+        if self.kind is TensorKind.PARAMETER:
+            return self.parent._parameters.get(self.leaf_name)
+        return self.parent._buffers.get(self.leaf_name)
+
+
+class FinalLayoutTensorRestorePlan(WeightRestorePlan):
+    """One-shot exact tensor rebinding validated before the first mutation."""
+
+    def __init__(
+        self,
+        lease: HostWeightLease,
+        replacements: tuple[_Replacement, ...],
+        validators: tuple[Callable[[], None], ...],
+        source_guard: Callable[[], None],
+    ) -> None:
+        self._lease = lease
+        self._replacements = replacements
+        self._validators = validators
+        self._source_guard = source_guard
+        self._committed = False
+
+    def commit(self) -> None:
+        if self._committed:
+            raise RuntimeError("final-layout tensor restore plan was already committed")
+        if self._lease.closed:
+            raise RuntimeError("cannot commit a final-layout tensor restore from a closed lease")
+        self._source_guard()
+
+        for replacement in self._replacements:
+            current = replacement.current_target()
+            if current is not replacement.target:
+                raise RuntimeError(f"restore target {replacement.name!r} changed after preflight")
+            if (
+                tuple(current.shape) != tuple(replacement.source.shape)
+                or current.dtype != replacement.source.dtype
+                or tuple(current.stride()) != tuple(replacement.source.stride())
+                or not replacement.source.is_contiguous()
+            ):
+                raise RuntimeError(f"restore target {replacement.name!r} changed layout after preflight")
+
+        # Mark first: assignment or model-validator failure makes the plan
+        # permanently non-retryable and the partially restored model disposable.
+        self._committed = True
+        for replacement in self._replacements:
+            if replacement.target.is_meta:
+                if replacement.kind is TensorKind.PARAMETER:
+                    replacement.parent._parameters[replacement.leaf_name] = nn.Parameter(
+                        replacement.source,
+                        requires_grad=replacement.target.requires_grad,
+                    )
+                else:
+                    replacement.parent._buffers[replacement.leaf_name] = replacement.source
+            else:
+                replacement.target.data = replacement.source
+        for validator in self._validators:
+            validator()
+
+
+class FinalLayoutTensorRestorer:
+    """Validate and plan restoration for one exact policy-defined lease."""
+
+    def __init__(self, context: FinalLayoutIdentityContext) -> None:
+        if not context.dit_names or any(not name for name in context.dit_names):
+            raise ValueError("final-layout restorer requires DiT component names")
+        self._context = context
+
+    @property
+    def schema(self) -> str:
+        return self._context.spec.producer.restorer_schema
+
+    def plan_restore(self, model: object, lease: HostWeightLease) -> FinalLayoutTensorRestorePlan:
+        if not isinstance(model, nn.Module):
+            raise TypeError("final-layout restoration requires an nn.Module pipeline")
+        if type(model) is not self._context.pipeline_type:
+            raise ValueError("restore model implementation differs from the exact identity context")
+        if lease.closed:
+            raise ValueError("cannot restore from a closed HostWeightLease")
+        if lease.identity.canonical_bytes != self._context.identity.canonical_bytes:
+            raise ValueError("lease semantic identity differs from the exact restore request")
+        if lease.manifest.identity.canonical_bytes != self._context.identity.canonical_bytes:
+            raise ValueError("lease manifest identity differs from the exact restore request")
+        if lease.identity.producer.restorer_schema != self.schema or lease.manifest.restorer_schema != self.schema:
+            raise ValueError("lease restorer schema is incompatible with the exact restore request")
+        self._context.ensure_sources_unchanged()
+
+        try:
+            dit_modules = tuple((name, model.get_submodule(name)) for name in self._context.dit_names)
+        except AttributeError as exc:
+            raise ValueError("one or more lease-owned DiT modules do not exist in the pipeline") from exc
+        if tuple(type(module) for _, module in dit_modules) != self._context.dit_types:
+            raise ValueError("restore DiT implementation differs from the exact identity context")
+        bindings = validate_final_layout_model_contract(
+            dit_modules,
+            expected_schema=self._context.spec.model_contract_schema,
+        )
+        targets = collect_final_layout_targets(
+            model,
+            dit_modules,
+            policy=self._context.policy,
+            require_materialized=False,
+        )
+        contract_digest = validate_final_layout_identity(self._context, targets)
+
+        manifest_entries = {entry.name: entry for entry in lease.manifest.tensors}
+        target_names = {target.name for target in targets}
+        if set(manifest_entries) != target_names or target_names != self._context.tensor_names:
+            missing = sorted(target_names - set(manifest_entries))
+            unexpected = sorted(set(manifest_entries) - target_names)
+            raise ValueError(
+                "lease tensor coverage differs from the structurally initialized DiT: "
+                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
+            )
+        self._context.policy.validate_format_metadata(
+            lease.manifest.format_metadata,
+            component_names=self._context.dit_names,
+            tensor_contract_digest=contract_digest,
+            tensor_count=len(manifest_entries),
+        )
+
+        replacements: list[_Replacement] = []
+        source_storages: dict[tuple[int, int], str] = {}
+        for target in targets:
+            entry = manifest_entries[target.name]
+            source = lease.tensors[target.name]
+            if entry.kind is not target.kind or entry.role != target.role:
+                raise ValueError(f"lease tensor kind or role differs for {target.name!r}")
+            if (
+                tuple(source.shape) != tuple(target.tensor.shape)
+                or source.dtype != target.tensor.dtype
+                or tuple(source.stride()) != tuple(target.tensor.stride())
+                or not source.is_contiguous()
+                or source.device.type != "cpu"
+            ):
+                raise ValueError(f"lease tensor layout differs for {target.name!r}")
+            if source.numel():
+                storage = source.untyped_storage()
+                storage_id = (storage.data_ptr(), storage.nbytes())
+                owner = source_storages.setdefault(storage_id, target.name)
+                if owner != target.name:
+                    raise ValueError(f"lease tensor {target.name!r} aliases {owner!r}")
+            parent_path, _, leaf_name = target.name.rpartition(".")
+            replacements.append(
+                _Replacement(
+                    name=target.name,
+                    parent=model.get_submodule(parent_path),
+                    leaf_name=leaf_name,
+                    target=target.tensor,
+                    source=source,
+                    kind=target.kind,
+                )
+            )
+
+        self._context.ensure_sources_unchanged()
+        return FinalLayoutTensorRestorePlan(
+            lease,
+            tuple(replacements),
+            tuple(binding.validator for binding in bindings),
+            self._context.ensure_sources_unchanged,
+        )
+
+
+__all__ = [
+    "FinalLayoutTensorRestorePlan",
+    "FinalLayoutTensorRestorer",
+]

vllm_omni/diffusion/model_loader/host_weights/source_identity.py

@@ -0,0 +1,351 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Canonical diffusion weight-source identity and one-load replacement guard."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import stat
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+import regex as re
+
+from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
+from vllm_omni.host_weight_runtime.identity import canonical_json
+
+from .contracts import (
+    FinalLayoutContractCode,
+    FinalLayoutContractError,
+    ImplementationIdentity,
+)
+
+_HASH_CHUNK_BYTES = 8 * 1024**2
+_IMMUTABLE_REVISION_RE = re.compile(r"(?:[0-9a-fA-F]{40}|[0-9a-fA-F]{64})\Z")
+
+
+class WeightSourceKind(str, Enum):
+    LOCAL_PATH = "local_path"
+    HUGGING_FACE_HUB = "hugging_face_hub"
+
+
+@dataclass(frozen=True)
+class PreparedWeightSource:
+    """Canonical source files and adapter ABI selected by the ordinary loader."""
+
+    model_or_path: str
+    subfolder: str | None
+    requested_revision: str | None
+    prefix: str
+    resolved_root: Path
+    weight_files: tuple[Path, ...]
+    use_safetensors: bool
+    checkpoint_adapter: ImplementationIdentity | None = None
+    source_kind: WeightSourceKind = WeightSourceKind.LOCAL_PATH
+
+    def __post_init__(self) -> None:
+        if not self.model_or_path or not isinstance(self.prefix, str):
+            raise ValueError("prepared source model and prefix must be valid strings")
+        if not isinstance(self.resolved_root, Path):
+            raise ValueError("prepared source root must use pathlib.Path")
+        if not isinstance(self.weight_files, tuple) or not self.weight_files:
+            raise ValueError("prepared source must contain an immutable tuple of weight files")
+        if any(not isinstance(path, Path) for path in self.weight_files):
+            raise ValueError("prepared source weight files must use pathlib.Path")
+        if not isinstance(self.use_safetensors, bool):
+            raise ValueError("prepared source safetensors mode must be a boolean")
+        if self.checkpoint_adapter is not None and not isinstance(
+            self.checkpoint_adapter,
+            ImplementationIdentity,
+        ):
+            raise ValueError("checkpoint adapter must use ImplementationIdentity")
+        if not isinstance(self.source_kind, WeightSourceKind):
+            raise ValueError("prepared source kind must use WeightSourceKind")
+
+
+@dataclass(frozen=True)
+class _FileSnapshot:
+    path: Path
+    relative_name: str
+    size: int
+    device: int
+    inode: int
+    mtime_ns: int
+    ctime_ns: int
+    symlink_target: str | None
+    content_id: str
+
+    def semantic_dict(self) -> dict[str, object]:
+        return {
+            "relative_name": self.relative_name,
+            "size": self.size,
+            "content_id": self.content_id,
+        }
+
+    def unchanged(self) -> bool:
+        try:
+            current = self.path.stat()
+            symlink_target = os.readlink(self.path) if self.path.is_symlink() else None
+        except OSError:
+            return False
+        return (
+            current.st_size == self.size
+            and current.st_dev == self.device
+            and current.st_ino == self.inode
+            and current.st_mtime_ns == self.mtime_ns
+            and current.st_ctime_ns == self.ctime_ns
+            and symlink_target == self.symlink_target
+        )
+
+
+@dataclass(frozen=True)
+class _SourceSnapshot:
+    semantic: CanonicalJson
+    revision: str
+    files: tuple[_FileSnapshot, ...]
+
+    def unchanged(self) -> bool:
+        return all(file.unchanged() for file in self.files)
+
+
+@dataclass(frozen=True)
+class FinalLayoutSourceContext:
+    """Resolved source identity plus process-local replacement observations."""
+
+    identity: WeightSourceIdentity
+    snapshots: tuple[_SourceSnapshot, ...]
+
+    def sources_unchanged(self) -> bool:
+        return all(snapshot.unchanged() for snapshot in self.snapshots)
+
+    def ensure_sources_unchanged(self) -> None:
+        if not self.sources_unchanged():
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.SOURCE_CHANGED,
+                "canonical source changed after final-layout identity resolution",
+            )
+
+
+def _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _hf_blob_content_id(
+    path: Path,
+    *,
+    source: PreparedWeightSource,
+    source_root: Path,
+) -> str | None:
+    if source.source_kind is not WeightSourceKind.HUGGING_FACE_HUB or not path.is_symlink():
+        return None
+    if Path(source.model_or_path).expanduser().exists():
+        return None
+
+    parts = source_root.parts
+    snapshot_indexes = [index for index, part in enumerate(parts[:-1]) if part == "snapshots"]
+    if not snapshot_indexes:
+        return None
+    snapshot_index = snapshot_indexes[-1]
+    revision = parts[snapshot_index + 1]
+    if _IMMUTABLE_REVISION_RE.fullmatch(revision) is None:
+        return None
+    repo_root = Path(*parts[:snapshot_index])
+    expected_repo_name = f"models--{source.model_or_path.replace('/', '--')}"
+    if repo_root.name != expected_repo_name:
+        return None
+    snapshot_root = repo_root / "snapshots" / revision
+    try:
+        path.relative_to(snapshot_root)
+    except ValueError:
+        return None
+    blobs_root = repo_root / "blobs"
+    try:
+        resolved_target = path.resolve(strict=True)
+        resolved_blobs_root = blobs_root.resolve(strict=True)
+    except OSError:
+        return None
+    if resolved_target.parent != resolved_blobs_root:
+        return None
+    blob_name = resolved_target.name
+    if _IMMUTABLE_REVISION_RE.fullmatch(blob_name) is None:
+        return None
+    return f"immutable-blob:{blob_name.lower()}"
+
+
+def _snapshot_revision(path: Path) -> str | None:
+    parts = path.absolute().parts
+    for index, part in enumerate(parts[:-1]):
+        if part != "snapshots":
+            continue
+        candidate = parts[index + 1]
+        if _IMMUTABLE_REVISION_RE.fullmatch(candidate) is not None:
+            return candidate.lower()
+    return None
+
+
+def _logical_model_id(value: str) -> str:
+    candidate = Path(value).expanduser()
+    try:
+        if candidate.exists():
+            return str(candidate.resolve())
+    except OSError:
+        pass
+    return value
+
+
+def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
+    root = source.resolved_root.expanduser().resolve()
+    if not root.is_dir():
+        raise ValueError(f"canonical source root is not a directory: {root}")
+    immutable_revision = _snapshot_revision(root)
+    if immutable_revision is None and source.requested_revision is not None:
+        requested = source.requested_revision.strip()
+        if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
+            immutable_revision = requested.lower()
+
+    files: list[_FileSnapshot] = []
+    for path in sorted(source.weight_files, key=lambda item: str(item)):
+        candidate = path.expanduser()
+        if not candidate.is_absolute():
+            candidate = root / candidate
+        candidate = candidate.absolute()
+        try:
+            relative_name = candidate.relative_to(root).as_posix()
+        except ValueError as exc:
+            raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
+        try:
+            current = candidate.stat()
+        except OSError as exc:
+            raise ValueError(f"cannot stat canonical weight file {candidate}") from exc
+        if not stat.S_ISREG(current.st_mode):
+            raise ValueError(f"canonical weight source is not a regular file: {candidate}")
+        symlink_target = os.readlink(candidate) if candidate.is_symlink() else None
+        content_id = _hf_blob_content_id(
+            candidate,
+            source=source,
+            source_root=root,
+        )
+        if content_id is None:
+            content_id = f"sha256:{_sha256_file(candidate)}"
+        files.append(
+            _FileSnapshot(
+                path=candidate,
+                relative_name=relative_name,
+                size=current.st_size,
+                device=current.st_dev,
+                inode=current.st_ino,
+                mtime_ns=current.st_mtime_ns,
+                ctime_ns=current.st_ctime_ns,
+                symlink_target=symlink_target,
+                content_id=content_id,
+            )
+        )
+
+    file_semantics = [file.semantic_dict() for file in files]
+    file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
+    revision = immutable_revision or f"content-{file_fingerprint}"
+    semantic = CanonicalJson.from_value(
+        {
+            "checkpoint_adapter": (
+                source.checkpoint_adapter.to_dict() if source.checkpoint_adapter is not None else None
+            ),
+            "files": file_semantics,
+            "model_or_path": _logical_model_id(source.model_or_path),
+            "prefix": source.prefix,
+            "resolved_revision": revision,
+            "source_kind": source.source_kind.value,
+            "subfolder": source.subfolder,
+            "use_safetensors": source.use_safetensors,
+        }
+    )
+    return _SourceSnapshot(semantic=semantic, revision=revision, files=tuple(files))
+
+
+def _resolve_target_sources(
+    prepared_sources: Sequence[PreparedWeightSource],
+    target_names: frozenset[str],
+) -> dict[str, PreparedWeightSource]:
+    """Bind every target to one source using deterministic longest-prefix wins."""
+    if not target_names:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
+            "final-layout source identity requires at least one owned tensor",
+        )
+
+    bindings: dict[str, PreparedWeightSource] = {}
+    for target_name in sorted(target_names):
+        matches = [source for source in prepared_sources if target_name.startswith(source.prefix)]
+        if not matches:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
+                f"no canonical weight source covers final-layout tensor {target_name!r}",
+            )
+        longest_prefix = max(len(source.prefix) for source in matches)
+        winners = [source for source in matches if len(source.prefix) == longest_prefix]
+        if len(winners) != 1:
+            candidates = sorted(
+                f"{source.model_or_path}:{source.subfolder or ''}:{source.prefix}" for source in winners
+            )
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.SOURCE_COVERAGE_INVALID,
+                f"multiple equally specific canonical sources cover {target_name!r}: {candidates}",
+            )
+        bindings[target_name] = winners[0]
+    return bindings
+
+
+def resolve_final_layout_source_identity(
+    prepared_sources: Sequence[PreparedWeightSource],
+    *,
+    model_id: str,
+    target_names: frozenset[str],
+) -> FinalLayoutSourceContext:
+    """Resolve exact source identity for files covering the owned tensors."""
+    bindings = _resolve_target_sources(prepared_sources, target_names)
+    selected_sources = frozenset(bindings.values())
+    snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources}
+    snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
+    source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
+    target_bindings = [
+        {
+            "source_fingerprint": hashlib.sha256(snapshot_by_source[source].semantic.encoded).hexdigest(),
+            "source_prefix": source.prefix,
+            "target_name": target_name,
+        }
+        for target_name, source in sorted(bindings.items())
+    ]
+    source_document = {
+        "sources": source_semantics,
+        "target_bindings": target_bindings,
+    }
+    source_fingerprint = hashlib.sha256(canonical_json(source_document)).hexdigest()
+    unique_revisions = tuple(sorted({snapshot.revision for snapshot in snapshots}))
+    aggregate_revision = (
+        unique_revisions[0]
+        if len(unique_revisions) == 1
+        else f"aggregate-{hashlib.sha256(canonical_json(unique_revisions)).hexdigest()}"
+    )
+    return FinalLayoutSourceContext(
+        identity=WeightSourceIdentity(
+            model_id=_logical_model_id(model_id),
+            revision=aggregate_revision,
+            fingerprint=source_fingerprint,
+            metadata=CanonicalJson.from_value(source_document),
+        ),
+        snapshots=snapshots,
+    )
+
+
+__all__ = [
+    "FinalLayoutSourceContext",
+    "PreparedWeightSource",
+    "WeightSourceKind",
+    "resolve_final_layout_source_identity",
+]

vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py

@@ -0,0 +1,287 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Representation-neutral ownership discovery for final-layout tensors."""
+
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass
+from itertools import chain
+
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
+from vllm_omni.host_weight_runtime import TensorKind
+from vllm_omni.host_weight_runtime.identity import canonical_json
+
+from .contracts import (
+    FinalLayoutContractCode,
+    FinalLayoutContractError,
+    FinalLayoutTensorPolicy,
+)
+
+
+@dataclass(frozen=True)
+class ModelRestoreBinding:
+    """Model-declared restore ABI and its post-commit validator."""
+
+    name: str
+    contract: FinalLayoutModelContract
+    validator: Callable[[], None]
+
+
+@dataclass(frozen=True)
+class RuntimeTensorTarget:
+    """One complete final-layout tensor owned by a DiT component."""
+
+    name: str
+    tensor: torch.Tensor
+    kind: TensorKind
+    role: str
+
+    @property
+    def nbytes(self) -> int:
+        return self.tensor.numel() * self.tensor.element_size()
+
+
+def validate_final_layout_model_contract(
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    *,
+    expected_schema: str,
+) -> tuple[ModelRestoreBinding, ...]:
+    """Require an explicit tensor-complete restore ABI from every DiT."""
+    if not dit_modules:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED,
+            "final-layout artifacts require at least one DiT component",
+        )
+
+    bindings: list[ModelRestoreBinding] = []
+    seen_names: set[str] = set()
+    for name, module in dit_modules:
+        if not name or name in seen_names:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+                f"DiT component name {name!r} is empty or duplicated",
+            )
+        seen_names.add(name)
+        contract = getattr(module, "host_weight_restore_contract", None)
+        if not isinstance(contract, FinalLayoutModelContract) or contract.schema != expected_schema:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED,
+                f"DiT component {name!r} does not declare restore schema {expected_schema!r}",
+            )
+        validator = getattr(module, "validate_restored_host_weights", None)
+        if not callable(validator):
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED,
+                f"DiT component {name!r} does not implement validate_restored_host_weights()",
+            )
+        bindings.append(ModelRestoreBinding(name, contract, validator))
+    return tuple(bindings)
+
+
+def _named_parameters_with_duplicates(module: nn.Module) -> Iterator[tuple[str, nn.Parameter]]:
+    try:
+        return module.named_parameters(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_parameters()
+
+
+def _named_buffers_with_duplicates(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
+    try:
+        return module.named_buffers(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_buffers()
+
+
+def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool:
+    parent_path, _, leaf_name = local_name.rpartition(".")
+    owner = module.get_submodule(parent_path)
+    return leaf_name not in owner._non_persistent_buffers_set
+
+
+def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
+    parent_path, _, leaf_name = runtime_name.rpartition(".")
+    parent = pipeline.get_submodule(parent_path)
+    tensor = parent._parameters.get(leaf_name)
+    if tensor is None:
+        tensor = parent._buffers.get(leaf_name)
+    return tensor
+
+
+def collect_final_layout_targets(
+    pipeline: nn.Module,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    *,
+    policy: FinalLayoutTensorPolicy,
+    require_materialized: bool,
+) -> tuple[RuntimeTensorTarget, ...]:
+    """Return complete, alias-free targets validated by one representation policy."""
+    validate_final_layout_model_contract(
+        dit_modules,
+        expected_schema=policy.spec.model_contract_schema,
+    )
+    records: dict[str, RuntimeTensorTarget] = {}
+    for dit_name, dit_module in dit_modules:
+        for local_name, tensor in _named_parameters_with_duplicates(dit_module):
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = RuntimeTensorTarget(
+                runtime_name,
+                tensor,
+                TensorKind.PARAMETER,
+                policy.tensor_role(runtime_name, tensor, TensorKind.PARAMETER),
+            )
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise FinalLayoutContractError(
+                    FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+        for local_name, tensor in _named_buffers_with_duplicates(dit_module):
+            if not _is_persistent_buffer(dit_module, local_name):
+                continue
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = RuntimeTensorTarget(
+                runtime_name,
+                tensor,
+                TensorKind.BUFFER,
+                policy.tensor_role(runtime_name, tensor, TensorKind.BUFFER),
+            )
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise FinalLayoutContractError(
+                    FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+
+    if not records:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+            "no DiT parameters or persistent buffers were discovered",
+        )
+
+    object_owners: dict[int, str] = {}
+    storage_owners: dict[tuple[int, int], str] = {}
+    for record in records.values():
+        tensor = record.tensor
+        if not record.role:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} has an empty semantic role",
+            )
+        try:
+            pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
+        except AttributeError as exc:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+                f"DiT tensor {record.name!r} is not owned by the pipeline",
+            ) from exc
+        if pipeline_tensor is not tensor:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS,
+                f"DiT tensor {record.name!r} does not resolve to the discovered object",
+            )
+
+        object_owner = object_owners.setdefault(id(tensor), record.name)
+        if object_owner != record.name:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} aliases tensor object {object_owner!r}",
+            )
+        if hasattr(tensor, "to_local"):
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} is a distributed tensor",
+            )
+        if tensor.layout != torch.strided:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} uses unsupported layout {tensor.layout}",
+            )
+        if tensor.device.type not in {"cpu", "meta"}:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} must be a CPU or meta tensor, got {tensor.device}",
+            )
+        if not tensor.is_contiguous():
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
+            )
+        policy.validate_target(record)
+
+        if not require_materialized:
+            continue
+        if tensor.device.type != "cpu" or tensor.is_meta:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
+            )
+        storage = tensor.untyped_storage()
+        if tensor.storage_offset() != 0 or storage.nbytes() != record.nbytes:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"{record.name!r} is a view into a larger storage",
+            )
+        if record.nbytes:
+            storage_id = (storage.data_ptr(), storage.nbytes())
+            storage_owner = storage_owners.setdefault(storage_id, record.name)
+            if storage_owner != record.name:
+                raise FinalLayoutContractError(
+                    FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                    f"{record.name!r} shares storage with {storage_owner!r}",
+                )
+
+    for pipeline_name, tensor in chain(
+        _named_parameters_with_duplicates(pipeline),
+        _named_buffers_with_duplicates(pipeline),
+    ):
+        owner = object_owners.get(id(tensor))
+        if owner is not None and owner != pipeline_name:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"cached tensor {owner!r} aliases pipeline tensor {pipeline_name!r}",
+            )
+        if not require_materialized or tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
+            continue
+        storage = tensor.untyped_storage()
+        pipeline_storage_owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
+        if pipeline_storage_owner is not None and pipeline_storage_owner != pipeline_name:
+            raise FinalLayoutContractError(
+                FinalLayoutContractCode.TENSOR_UNSUPPORTED,
+                f"cached tensor {pipeline_storage_owner!r} shares storage with pipeline tensor {pipeline_name!r}",
+            )
+
+    ordered = tuple(records[name] for name in sorted(records))
+    policy.validate_collection(ordered)
+    return ordered
+
+
+def tensor_contract_sha256(records: Sequence[RuntimeTensorTarget]) -> str:
+    """Hash exact structural ownership and representation-specific tensor roles."""
+    contract = [
+        {
+            "dtype": str(record.tensor.dtype),
+            "kind": record.kind.value,
+            "name": record.name,
+            "role": record.role,
+            "shape": list(record.tensor.shape),
+            "stride": list(record.tensor.stride()),
+        }
+        for record in records
+    ]
+    return hashlib.sha256(canonical_json(contract)).hexdigest()
+
+
+__all__ = [
+    "ModelRestoreBinding",
+    "RuntimeTensorTarget",
+    "collect_final_layout_targets",
+    "tensor_contract_sha256",
+    "validate_final_layout_model_contract",
+]

vllm_omni/diffusion/models/host_weight_contract.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Model-owned declaration for complete final-layout tensor restoration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA = "vllm-omni.diffusion.final-layout-tensors-v1"
+
+
+@dataclass(frozen=True)
+class FinalLayoutModelContract:
+    """Versioned model ABI for reconstruction from exact final-layout tensors."""
+
+    implementation_id: str
+    version: str
+    schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+
+    def __post_init__(self) -> None:
+        for name, value in (
+            ("schema", self.schema),
+            ("implementation_id", self.implementation_id),
+            ("version", self.version),
+        ):
+            if not isinstance(value, str) or not value:
+                raise ValueError(f"final-layout model contract {name} must not be empty")
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "schema": self.schema,
+            "implementation_id": self.implementation_id,
+            "version": self.version,
+        }
+
+
+__all__ = [
+    "FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA",
+    "FinalLayoutModelContract",
+]

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 packed-token audio/video DiT for vLLM-Omni.
 
 vLLM tensor parallel linears and the unified attention layer provide TP and
@@ -43,6 +44,7 @@ from vllm_omni.diffusion.layers.activation import SiluAndMul
 from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -889,6 +891,15 @@ class MiniMaxH3SPGather(nn.Module):
 
 
 class MiniMaxH3DiTModel(nn.Module):
+    # Loading is tensor-complete: constructor state plus final-layout
+    # parameters and persistent buffers is sufficient to reconstruct a ready
+    # inference model. The model-specific validator below checks the preserved
+    # FP32 portion after a lease-backed restore commits.
+    host_weight_restore_contract = FinalLayoutModelContract(
+        implementation_id="minimax-h3-dit",
+        version="1",
+    )
+
     _cache_dit_adapter_config = CacheDiTAdapterConfig(
         block_forward_patterns={"blocks": ForwardPattern.Pattern_3},
         # H3 is CFG-distilled and performs one transformer forward per step.
@@ -1064,6 +1075,10 @@ class MiniMaxH3DiTModel(nn.Module):
             if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE:
                 raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.")
 
+    def validate_restored_host_weights(self) -> None:
+        """Validate mixed-precision invariants after lease-backed restore."""
+        self.post_load_weights()
+
     def load_weights(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],