[Diffusion] Integrate no-AllGather DLO with Host Weight Runtime (#6486)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -20,7 +20,7 @@ implemented but has less end-to-end coverage than the primary path.
 Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
 
 | Feature | DLO + AllGather | DLO without AllGather |
-|---|---|---|
+| --- | --- | --- |
 | **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
 | **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
 | **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
@@ -169,10 +169,167 @@ This mode means:
   copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
+### Final-layout Host Weight Runtime consumer
+
+> **Status:** PR2 implementation. The representation-independent identity,
+> producer, and restoration contracts landed in
+> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
+> wires the final-layout BF16 path into eligible no-AllGather DLO.
+
+Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
+only to no-AllGather DLO. Enable it with
+`--host-weight-runtime-mode preferred` (or `required`) and
+`--host-weight-runtime-root <node-local-root>`. Registration of shared mappings
+and direct asynchronous H2D remain a separate PR3 transport change.
+
+The modes express operator fallback policy, not different artifact formats:
+
+- `preferred` consumes an exact hit, but on a miss it allows canonical loading
+  followed by post-load publication. It is the normal population path.
+- `required` consumes the same exact artifact but fails startup on a miss or
+  unusable artifact. It never invokes canonical DiT fallback or post-load
+  publication and therefore cannot populate an empty store.
+
+PR2 has no separate prewarm command. Operators populate one matching producer
+cohort per node-local storage domain in `preferred` mode, then restart the same
+model revision and parallel layout with `required`. TP coordinates have distinct
+identities, while equivalent DP replicas share them.
+
+The current producer/consumer boundary is the model-declared final-layout BF16
+contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
+and TP2 rank identities plus SP layout identities; online quantization, HSDP,
+LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
+`disabled`, DLO-disabled, and DLO AllGather configurations stop before source
+identity or store construction and retain the existing checkpoint-mmap or
+ordinary-loader path.
+
+Eligibility is decided before constructing HWR, resolving canonical sources,
+hashing identity inputs, probing a filesystem, or emitting an HWR observer
+event:
+
+```text
+DLO disabled / HWR disabled / DLO AllGather
+  -> zero final-layout HWR interaction
+  -> preserve current checkpoint-mmap or ordinary-loader behavior
+```
+
+When final-layout HWR is explicitly enabled for no-AllGather DLO, runtime
+outcomes are authoritative:
+
+- `LOCAL_HIT`: plan and commit an exact restoration, then transfer the lease
+  transactionally to DLO.
+- `CANONICAL_FALLBACK`: bypass checkpoint mmap, canonically materialize and
+  finalize the DiT, then call `publish_after_load()` for a future startup. The
+  current startup retains its canonical tensors.
+- `FAILED`: fail startup. Preferred mode does not reinterpret nonretryable
+  identity, configuration, or compatibility failures as misses.
+- Required mode cannot bootstrap an empty store through a `POST_LOAD_ONLY`
+  producer.
+
+Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
+not selected.
+
+#### Pre-service transaction boundary
+
+The startup transaction does not end at restore commit. A lease-backed model is
+disposable until backend setup and initial prefetch complete:
+
+```text
+UNRESOLVED
+  -> LEASE_OWNED_BY_LOADER
+  -> RESTORE_PLANNED
+  -> COMMIT_STARTED
+  -> CARRIER_OWNED
+  -> BACKEND_OWNED
+  -> BACKEND_READY
+  -> IN_SERVICE
+```
+
+Failure before `COMMIT_STARTED` closes the loader-owned lease. Preferred mode
+may canonically load using the untouched skeleton; required mode fails.
+
+Failure from `COMMIT_STARTED` through `BACKEND_READY` must:
+
+1. synchronize or otherwise quiesce partial backend and initial-prefetch work;
+2. release hook, staging, restore-plan, and model references;
+3. close the lease through its current owner;
+4. discard the restored model; and
+5. in preferred mode, construct a fresh canonical model while bypassing HWR
+   lookup, HWR publication, and checkpoint mmap for that recovery attempt.
+
+Required mode fails instead of constructing the fresh fallback. Once the model
+enters service, failures follow normal runtime handling rather than startup
+fallback.
+
+#### Finalization phases
+
+Cold canonical loading separates byte-changing work from shared runtime-state
+finalization:
+
+```text
+Cold-only byte-changing:
+  casting, reordering, packing, quantization, generated scales,
+  model-specific weight transforms, and calibration that mutates
+  parameters or persistent buffers
+
+Shared cold/warm non-byte:
+  validation, hook installation, eval state, bookkeeping,
+  and non-persistent runtime state
+```
+
+The warm path must snapshot restored tensor bytes and backing pointers before
+shared finalization and prove both are unchanged afterward.
+
+#### Lease ownership
+
+Lease transfer is single-owner and single-take:
+
+```text
+loader owns
+  -> restore plan borrows
+  -> commit succeeds
+  -> runner carrier takes once
+  -> backend takes once before asynchronous work
+  -> backend drains pending H2D work
+  -> backend releases hooks, staging, and model references
+  -> backend closes lease
+```
+
+The runner carrier is process-local, rejects serialization, and rejects a
+duplicate `take()`. If the carrier still owns the lease, runner cleanup closes
+it. Once the backend takes ownership, only backend abort or teardown may drain
+work and close it.
+
+The implementation keeps the final-layout lease through backend setup and
+initial prefetch. The backend uses the existing two bounded rank-local staging
+slots for the immutable HWR tensors, and its transactional `enable()` cleanup
+removes partial hooks and closes the lease before reporting a setup failure.
+Preferred mode then uses the runner's fresh canonical retry; required mode
+propagates the failure.
+
+#### PR2 promotion gates
+
+- Warm hit performs zero ordinary DiT materialization and zero producer calls.
+- Shared warm finalization changes neither restored bytes nor backing pointers.
+- Preferred `FAILED` outcomes do not silently fall back.
+- Planning failure leaves the skeleton untouched.
+- Any pre-service failure after commit begins discards the restored model;
+  preferred mode uses a fresh canonical model and required mode fails.
+- Mixed components load normally and no required tensor remains on `meta`.
+- Disabled mode and AllGather emit zero final-layout HWR interaction.
+- The lease carrier rejects duplicate take and serialization.
+- Backend setup or prefetch failure drains asynchronous work before lease close.
+- Compatible checkpoint mmap remains an unchanged benchmark and control path.
+- Prewarm uses one matching producer cohort per TP/SP topology and storage
+  domain.
+- Benchmarks compare ordinary loading, cold publication, and warm staging with
+  output parity, startup latency, aggregate PSS, page-cache state, HBM, and H2D
+  payload.
+
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
-|---|---|---|
+| --- | --- | --- |
 | **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
 | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
 | **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
@@ -185,7 +342,8 @@ This mode means:
   the SP group becomes DLO's sharding group in AllGather mode.
 - **DP + TP/SP without AllGather:** standard model loading defines the
   rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
-  collective.
+  collective. When the model declares the final-layout contract, HWR keys the
+  reusable artifact by TP rank/size and SP layout.
 - **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
   DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
   materialization and synchronization.
@@ -211,10 +369,11 @@ reconstructs those tensors with their recorded layouts. Other online methods
 must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+The Host Weight Runtime representation and publication contracts are merged;
+see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
+no-AllGather consumer above is the PR2 implementation; registration and direct
+asynchronous H2D remain a separate transport follow-up.
 
 ## Validation coverage
 
@@ -243,7 +402,7 @@ VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
 or shared-mmap host-memory savings.
 
 | Configuration | Result | Warm E2E | Peak device memory | Host PSS |
-|---|---:|---:|---:|---:|
+| --- | ---: | ---: | ---: | ---: |
 | DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
 | DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
 | DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
@@ -275,7 +434,7 @@ pipeline components, so each worker should be compared with the same worker in
 the other storage mode.
 
 | Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
-|---|---:|---:|---:|---:|---:|
+| --- | ---: | ---: | ---: | ---: | ---: |
 | DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
 | DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
 | **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -63,6 +63,8 @@ omni = Omni(
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
+| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
 
 ## Host-weight loading
 
@@ -104,6 +106,88 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+### Final-layout Host Weight Runtime
+
+HWR is an opt-in startup optimization for models that declare the final-layout
+BF16 restore contract. Use it only with no-AllGather DLO:
+
+```bash
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
+  --host-weight-runtime-mode preferred \
+  --host-weight-runtime-root /var/cache/vllm-omni/hwr
+```
+
+#### Choosing a mode
+
+| Mode | Exact local artifact hit | Miss or recoverable artifact/store problem | Intended use |
+| --- | --- | --- | --- |
+| `disabled` | HWR is not consulted | Use the existing checkpoint-mmap or ordinary-loader path | Default compatibility path |
+| `preferred` | Restore the final-layout artifact | Load canonically, serve with those tensors, and attempt to publish an artifact for the next startup | Normal deployment and store population |
+| `required` | Restore the final-layout artifact | Fail startup without canonical DiT fallback or publication | Enforce a pre-populated store in controlled rollouts or CI |
+
+Both enabled modes still fail on non-retryable configuration, identity, or
+compatibility errors. `preferred` is a fallback policy for a cache miss or a
+recoverable cache problem; it does not hide an invalid deployment.
+
+#### Populating a store for `required` mode
+
+`required` is deliberately consume-only. PR2 does not include a separate
+prewarm command, so populate each node-local storage domain with one matching
+`preferred` producer cohort:
+
+1. Choose a persistent, writable root visible to every diffusion worker in the
+   storage domain. Do not use a process-private temporary directory.
+2. Start the deployment with `preferred`, using the exact model revision,
+   dtype, TP size, SP configuration, and other weight-layout settings intended
+   for serving. No inference request is needed; publication happens during
+   model startup. Wait for the engine to become healthy, then shut it down
+   cleanly.
+3. Restart with the same arguments and root, changing only the mode to
+   `required`. A successful startup proves that every worker acquired a valid
+   artifact; a miss, corrupt artifact, or incompatible identity fails startup.
+4. Repeat the population start on every node or storage domain because this
+   store is node-local.
+
+For example:
+
+```bash
+# First startup: canonically load and populate the exact artifacts.
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
+  --host-weight-runtime-mode preferred \
+  --host-weight-runtime-root /var/cache/vllm-omni/hwr
+
+# After a healthy startup and clean shutdown, enforce cache hits.
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
+  --host-weight-runtime-mode required \
+  --host-weight-runtime-root /var/cache/vllm-omni/hwr
+```
+
+Include the same TP and SP layout flags in both commands. DP rank and DP size
+are excluded from artifact identity, so the population and serving DP sizes may
+differ and equivalent DP replicas share artifacts. TP rank is included, so a
+TP-N deployment normally needs N rank-specific artifacts in each storage
+domain; launching the matching TP cohort creates that set. If the model
+revision or layout changes, run `preferred` again for the new identity before
+returning to `required`.
+
+On a cold start, the canonical loader remains authoritative and publishes a
+validated final-layout artifact for later workers. A warm start restores the
+DiT final tensors without ordinary DiT materialization, then DLO streams them
+through the same two bounded host staging slots. The artifact identity includes
+the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP
+layouts do not alias one another. Publication failure remains separate from a
+`preferred` serving startup; the next `required` startup provides the explicit
+artifact-availability check.
+
+When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
+not resolve HWR sources or construct its store.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -135,7 +219,9 @@ must enter each collective.
 - Direct checkpoint mmap currently requires TP1. TP greater than one is
   outside the Phase A shared-mmap support scope and falls back before model
   mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  but it provides no shared-mmap host-memory benefit or guarantee.
+  and eligible no-AllGather configurations may use HWR final-layout artifacts,
+  but direct checkpoint mmap provides no shared-mmap host-memory guarantee for
+  TP greater than one.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
 - Per-tensor online FP8 linears use the ordinary loader and can run with either
@@ -147,10 +233,10 @@ must enter each collective.
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
+Sharing quantized or otherwise unvalidated transformed layouts through a
+normalized HWR producer is a follow-up design in
 [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+the current BF16 final-layout path.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,17 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Tests for the DiffusersPipelineLoader.
 """
 
 import json
+from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
+from safetensors.torch import save_file
 from vllm.config.load import LoadConfig
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
@@ -23,6 +25,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     TensorBinding,
 )
 from vllm_omni.diffusion.models.helios import HeliosPipeline
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -73,6 +76,73 @@ class _DummyPipelineModel(nn.Module):
         return loaded
 
 
+class _HWRTransformer(nn.Module):
+    host_weight_restore_contract = FinalLayoutModelContract(
+        implementation_id="test-hwr-transformer",
+        version="1",
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))
+
+    def validate_restored_host_weights(self):
+        assert self.weight.dtype is torch.bfloat16
+
+
+class _HWRPipeline(nn.Module):
+    def __init__(self, source_root: Path):
+        super().__init__()
+        self.transformer = _HWRTransformer()
+        self.load_count = 0
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=str(source_root),
+                subfolder=None,
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=False,
+            )
+        ]
+
+    def load_weights(self, weights):
+        loaded: set[str] = set()
+        params = dict(self.named_parameters())
+        for name, tensor in weights:
+            if name in params:
+                params[name].data.copy_(tensor.to(dtype=params[name].dtype))
+                loaded.add(name)
+        self.load_count += 1
+        return loaded
+
+
+def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace:
+    return SimpleNamespace(
+        model=str(model),
+        dtype=torch.bfloat16,
+        host_weight_runtime_mode=mode,
+        host_weight_runtime_root=str(root),
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        lora_path=None,
+        quantization_config=None,
+        diffusion_attention_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            data_parallel_size=1,
+            sequence_parallel_size=1,
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            cfg_parallel_size=1,
+            enable_expert_parallel=False,
+            ulysses_degree=1,
+            ring_degree=1,
+            allgather_degree=1,
+            ulysses_mode="strict",
+        ),
+    )
+
+
 def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
     od_config = SimpleNamespace(
         dtype=torch.float32,
@@ -92,6 +162,81 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
     return loader
 
 
+def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
+    tmp_path: Path,
+    monkeypatch,
+):
+    canonical_root = tmp_path / "canonical"
+    canonical_root.mkdir()
+    save_file(
+        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+        str(canonical_root / "model.safetensors"),
+    )
+    store_root = tmp_path / "hwr-store"
+
+    def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
+        loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
+        pipeline = _HWRPipeline(canonical_root)
+        monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+        return loader, pipeline
+
+    cold_loader, cold_model = make_loader()
+    cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+    assert cold is cold_model
+    assert cold_model.load_count == 1
+    assert cold_loader._hwr_state is not None
+
+    warm_loader, warm_model = make_loader()
+    monkeypatch.setattr(
+        warm_loader,
+        "_process_weights_after_loading",
+        lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"),
+    )
+    warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    assert warm is warm_model
+    assert warm_model.load_count == 0
+    assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight)
+    from vllm_omni.diffusion.offloader.startup import take_offload_startup_state
+
+    startup_state = take_offload_startup_state(warm)
+    assert startup_state is not None
+    warm_plan = startup_state.host_weight_plan
+    assert warm_plan is not None
+    assert warm_plan.lease_carrier is not None
+    warm_plan.lease_carrier.close()
+
+
+def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
+    from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module
+
+    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store"))
+    models: list[_DummyPipelineModel] = []
+
+    def init_model(*args, **kwargs):
+        del args, kwargs
+        model = _DummyPipelineModel(source_prefix="transformer.")
+        models.append(model)
+        return model
+
+    def commit_error(*args, **kwargs):
+        del args, kwargs
+        raise loader_module._HWRCommitError("restore commit failed")
+
+    monkeypatch.setattr(loader, "_init_from_load_format", init_model)
+    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ())
+    monkeypatch.setattr(loader, "_resolve_hwr", commit_error)
+    monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None)
+    monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None)
+    monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None)
+
+    recovered = loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    assert len(models) == 2
+    assert recovered is models[1]
+    assert loader.take_host_weight_plan() is None
+
+
 def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
@@ -103,6 +248,64 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     )
 
 
+@pytest.mark.parametrize(
+    ("dist_offload", "use_allgather"),
+    [
+        (False, False),
+        (True, True),
+    ],
+)
+def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
+    monkeypatch,
+    tmp_path,
+    dist_offload,
+    use_allgather,
+):
+    """Disabled and AllGather paths must never construct or probe HWR."""
+    from vllm_omni.host_weight_runtime import HostWeightRuntime
+
+    root = tmp_path / "must-not-be-touched"
+    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
+
+    def unexpected_store_construction(*args, **kwargs):
+        raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}")
+
+    monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction)
+    assert (
+        loader._resolve_hwr(
+            model,
+            modules,
+            dist_offload=dist_offload,
+            use_allgather=use_allgather,
+            load_format="default",
+            sources=tuple(model.weights_sources),
+        )
+        is None
+    )
+    assert not root.exists()
+
+
+def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path):
+    loader = DiffusersPipelineLoader(
+        LoadConfig(),
+        _hwr_config("dummy-model", tmp_path / "store", mode="required"),
+    )
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
+
+    with pytest.raises(ValueError, match="restore contract"):
+        loader._resolve_hwr(
+            model,
+            modules,
+            dist_offload=True,
+            use_allgather=False,
+            load_format="default",
+            sources=tuple(model.weights_sources),
+        )
+
+
 @pytest.mark.parametrize("offline", [False, True])
 def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -726,6 +726,36 @@ def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) ->
     assert identity_v2.key != base_identity.key
 
 
+@pytest.mark.parametrize(
+    "parallel",
+    [
+        FinalLayoutParallelIdentity(),
+        FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0),
+        FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1),
+        FinalLayoutParallelIdentity(sequence_parallel_size=2, ulysses_degree=2),
+        FinalLayoutParallelIdentity(sequence_parallel_size=2, ring_degree=2),
+        FinalLayoutParallelIdentity(
+            tensor_parallel_size=2,
+            tensor_parallel_rank=1,
+            sequence_parallel_size=2,
+            ulysses_degree=2,
+        ),
+    ],
+)
+def test_parallel_identity_matrix_is_exact_for_tp_and_sp(tmp_path: Path, parallel: FinalLayoutParallelIdentity) -> None:
+    """TP rank/size and SP backend are identity coordinates for HWR artifacts."""
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    context = _identity(model, source, request=_request(parallel=parallel))
+    layout = context.identity.layout
+
+    assert layout.tensor_parallel_size == parallel.tensor_parallel_size
+    assert layout.tensor_parallel_rank == parallel.tensor_parallel_rank
+    assert layout.sequence_parallel_size == parallel.sequence_parallel_size
+    assert layout.sequence_parallel_backend == parallel.sequence_parallel_backend
+    assert "data_parallel" not in context.identity.layout.metadata.to_value().get("parallel", {})
+
+
 @pytest.mark.parametrize(
     ("semantic_request", "message"),
     [

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
+from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -695,6 +696,67 @@ class _SingleBlockModel(nn.Module):
         self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
 
 
+class _HWRPipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=3)
+
+
+class _FakeHostWeightLease:
+    def __init__(self, resolution_id: str):
+        self.closed = False
+        self.provenance = SimpleNamespace(resolution_id=resolution_id)
+
+    def close(self):
+        self.closed = True
+
+
+class _FakeLeaseCarrier:
+    def __init__(self, lease: _FakeHostWeightLease):
+        self.lease = lease
+        self.taken = False
+
+    @property
+    def closed(self):
+        return self.taken or self.lease.closed
+
+    def take(self):
+        if self.taken:
+            raise RuntimeError("already taken")
+        self.taken = True
+        return self.lease
+
+    def close(self):
+        if not self.taken:
+            self.lease.close()
+
+
+def _fake_hwr_plan(resolution_id: str = "hwr-test"):
+    lease = _FakeHostWeightLease(resolution_id)
+    carrier = _FakeLeaseCarrier(lease)
+    plan = HostWeightPlan(
+        backing_kind="host_weight_runtime",
+        bindings={},
+        lease_carrier=carrier,
+    )
+    return plan, carrier, lease
+
+
+def _hwr_backend():
+    plan, carrier, lease = _fake_hwr_plan()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+        host_weight_plan=plan,
+    )
+    return _HWRPipeline(), backend, carrier, lease
+
+
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -838,6 +900,47 @@ class TestMmapWeightLoading:
         assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
         backend.disable()
 
+    def test_hwr_carrier_is_taken_and_released_after_bounded_staging(self, patched_offload_runtime):
+        """A warm HWR plan uses the existing two-slot rank-local transport."""
+        pipeline, backend, carrier, lease = _hwr_backend()
+        model_bytes = sum(param.numel() * param.element_size() for param in pipeline.transformer.parameters())
+
+        backend.enable(pipeline)
+
+        assert carrier.taken
+        assert backend._using_rank_local_mmap
+        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+        staging_bytes = sum(
+            buffer.numel() * buffer.element_size()
+            for buffer in backend._all_hook_groups[0][0].cpu_staging_buffers[0].values()
+        )
+        assert staging_bytes * 2 < model_bytes
+        assert not lease.closed
+
+        backend.disable()
+        assert lease.closed
+
+    def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
+        self,
+        monkeypatch,
+        patched_offload_runtime,
+    ):
+        pipeline, backend, carrier, lease = _hwr_backend()
+        monkeypatch.setattr(
+            backend,
+            "_allocate_shared_buffers",
+            lambda hooks: (_ for _ in ()).throw(RuntimeError("device buffer allocation failed")),
+        )
+
+        with pytest.raises(RuntimeError, match="device buffer allocation failed"):
+            backend.enable(pipeline)
+
+        assert carrier.taken
+        assert lease.closed
+        assert not backend.enabled
+        assert not backend._blocks
+        assert not backend._all_hook_groups
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -1504,6 +1607,48 @@ class TestMmapValidation:
 class TestConfigValidation:
     """Tests for configuration validation in OffloadConfig / execute_request."""
 
+    @pytest.mark.parametrize(
+        ("dp", "sp", "tp", "use_allgather", "expected_dlo_group"),
+        [
+            (1, 1, 1, True, 1),
+            (2, 1, 1, True, 2),
+            (1, 4, 1, True, 4),
+            (2, 2, 2, True, 2),
+            (2, 2, 2, False, 1),
+            (1, 4, 2, False, 1),
+        ],
+    )
+    def test_dlo_group_selection_covers_dp_sp_tp_matrix(
+        self,
+        dp,
+        sp,
+        tp,
+        use_allgather,
+        expected_dlo_group,
+    ):
+        """DLO sharding follows DP, falls back to SP, and never uses TP."""
+
+        config = OffloadConfig.from_od_config(
+            SimpleNamespace(
+                enable_cpu_offload=False,
+                enable_layerwise_offload=False,
+                enable_distributed_layerwise_offload=True,
+                dlo_use_allgather=use_allgather,
+                dlo_resident_layers=0,
+                pin_cpu_memory=False,
+                parallel_config=SimpleNamespace(
+                    data_parallel_size=dp,
+                    use_hsdp=False,
+                    hsdp_shard_size=-1,
+                    hsdp_replicate_size=1,
+                    sequence_parallel_size=sp,
+                    tensor_parallel_size=tp,
+                ),
+            )
+        )
+        assert config.dp_size == expected_dlo_group
+        assert config.dlo_use_allgather is use_allgather
+
     def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
         import vllm_omni.diffusion.offloader as offloader_module
 
@@ -1538,6 +1683,53 @@ class TestConfigValidation:
                 host_weight_plan=plan,
             )
 
+    def test_startup_recovery_stays_inside_offloader_boundary(self, monkeypatch):
+        import vllm_omni.diffusion.offloader as offloader_module
+
+        class FailingBackend:
+            def __init__(self):
+                self.disabled = False
+
+            def enable(self, pipeline):
+                del pipeline
+                raise RuntimeError("initial prefetch failed")
+
+            def disable(self):
+                self.disabled = True
+
+        warm_pipeline = nn.Module()
+        canonical_pipeline = nn.Module()
+        plan, carrier, _lease = _fake_hwr_plan("hwr-recovery")
+        failing_backend = FailingBackend()
+        calls = []
+
+        def fake_backend(od_config, device, host_weight_plan):
+            del od_config, device
+            calls.append(host_weight_plan)
+            return failing_backend if host_weight_plan is plan else None
+
+        attach_offload_startup_state(
+            warm_pipeline,
+            OffloadStartupState(
+                host_weight_plan=plan,
+                fresh_model_loader=lambda: canonical_pipeline,
+                allow_fresh_retry=True,
+            ),
+        )
+        monkeypatch.setattr(offloader_module, "get_offload_backend", fake_backend)
+
+        recovered, backend = offloader_module.enable_offload_backend(
+            SimpleNamespace(),
+            warm_pipeline,
+            device=torch.device("cpu"),
+        )
+
+        assert recovered is canonical_pipeline
+        assert backend is None
+        assert failing_backend.disabled
+        assert carrier.closed
+        assert calls == [plan, None]
+
     def test_hsdp_with_allgather_rejected(self):
         """HSDP + DLO + AllGather should raise ValueError (double sharding)."""
         from vllm_omni.diffusion.offloader.base import OffloadConfig

tests/diffusion/test_diffusion_model_runner.py

@@ -825,9 +825,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
             del kwargs
             return SimpleNamespace(transformer=torch.nn.Identity())
 
-        def take_host_weight_plan(self):
-            return None
-
     class _DummyMemoryProfiler:
         consumed_memory = 0
 
@@ -869,8 +866,8 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
     monkeypatch.setattr(
         model_runner_module,
-        "get_offload_backend",
-        lambda od_config, device, host_weight_plan: None,
+        "enable_offload_backend",
+        lambda od_config, pipeline, device: (pipeline, None),
     )
     monkeypatch.setattr(
         model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
@@ -883,8 +880,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert dummy_cache_backend.enabled is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
 def test_set_forward_context_enters_vllm_config_contexts(monkeypatch):
     """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions."""
     import vllm.config.vllm as vllm_config_module

tests/diffusion/test_diffusion_step_pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for step-level diffusion execution across runner / worker / executor / engine."""
 
 import contextlib
@@ -879,7 +879,11 @@ class TestRunner:
 
         monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
         monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
-        monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+        monkeypatch.setattr(
+            model_runner_module,
+            "enable_offload_backend",
+            lambda _config, pipeline, **_kwargs: (pipeline, None),
+        )
         monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
 
         with pytest.raises(ValueError, match="RequestOnlyPipeline"):

tests/diffusion/test_diffusion_streaming_output.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Integration tests for diffusion pipeline streaming output mode.
 
 Covers:
@@ -544,7 +544,11 @@ class TestSupportedPipelines:
 
         monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader)
         monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler)
-        monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+        monkeypatch.setattr(
+            model_runner_module,
+            "enable_offload_backend",
+            lambda _config, pipeline, **_kwargs: (pipeline, None),
+        )
         monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
 
         with pytest.raises(ValueError, match="NoStepPipeline"):

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -362,6 +362,35 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert engine_args["dlo_resident_layers"] == 20
 
 
+def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "MiniMaxAI/MiniMax-H3",
+            "--omni",
+            "--enable-distributed-layerwise-offload",
+            "--dlo-no-use-allgather",
+            "--host-weight-runtime-mode",
+            "preferred",
+            "--host-weight-runtime-root",
+            "/var/cache/vllm-omni/hwr",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    engine_args = stage_cfg["engine_args"]
+
+    assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
+    assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert engine_args["host_weight_runtime_mode"] == "preferred"
+    assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+
+
 def test_serve_cli_accepts_diffusion_compile_controls():
     """Ensure both compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()

tests/host_weight_runtime/test_filesystem_store.py

@@ -389,6 +389,24 @@ def test_build_lookup_share_one_published_backing_and_lease_blocks_cleanup(tmp_p
     assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS
 
 
+def test_publication_hardens_artifact_after_atomic_rename(tmp_path: Path) -> None:
+    identity = _identity()
+    store = _make_store(tmp_path / "store")
+
+    built = store.get_or_build(
+        BuildRequest(identity),
+        FakeProducer(identity),
+        validation=ValidationLevel.FULL_CHECKSUM,
+        deadline=time.monotonic() + 10,
+    )
+
+    assert built.status is StoreStatus.BUILT
+    assert built.lease is not None
+    built.lease.close()
+    artifact_dir = store.artifacts_dir / identity.key
+    assert artifact_dir.stat().st_mode & 0o777 == 0o555
+
+
 def test_concurrent_lease_close_waits_for_resource_teardown(tmp_path: Path) -> None:
     identity = _identity()
     store = _make_store(tmp_path / "store")

tests/host_weight_runtime/test_runtime_resolution.py

@@ -17,6 +17,7 @@ from vllm_omni.host_weight_runtime import (
     FailureCode,
     HostWeightFailure,
     HostWeightLease,
+    HostWeightLeaseCarrier,
     HostWeightRuntime,
     HostWeightRuntimeConfig,
     IntegrityPolicy,
@@ -166,6 +167,28 @@ def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp
     assert all(report.resolution_id for report in reports)
 
 
+def test_lease_carrier_is_single_take_and_process_local(tmp_path: Path) -> None:
+    identity = _identity()
+    producer = CountingProducer(identity)
+    runtime = HostWeightRuntime.from_config(
+        HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store"))
+    )
+    resolution = runtime.resolve(identity, producer=producer)
+    assert resolution.lease is not None
+
+    carrier = HostWeightLeaseCarrier(resolution.lease)
+    with pytest.raises(TypeError, match="cannot be serialized"):
+        carrier.__reduce_ex__(4)
+
+    lease = carrier.take()
+    assert lease is resolution.lease
+    with pytest.raises(RuntimeError, match="already taken"):
+        carrier.take()
+    carrier.close()  # ownership has moved; the carrier must not close twice.
+    assert not lease.closed
+    lease.close()
+
+
 @pytest.mark.parametrize(
     ("mode", "expected"),
     [

vllm_omni/config/omni_config.py

@@ -696,6 +696,8 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
+    host_weight_runtime_root: str | None = None
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -753,6 +755,7 @@ class _DiffusionConfigProjection:
             TransformerConfig,
             build_attention_config,
             parse_kv_cache_skip_selector,
+            validate_host_weight_runtime_options,
         )
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
@@ -830,6 +833,11 @@ class _DiffusionConfigProjection:
         elif self.max_cpu_loras < 1:
             raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
 
+        validate_host_weight_runtime_options(
+            mode=self.host_weight_runtime_mode,
+            root=self.host_weight_runtime_root,
+        )
+
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(
                 "diffusers_load_kwargs and diffusers_call_kwargs are only "

vllm_omni/config/stage_config.py

@@ -386,6 +386,8 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    host_weight_runtime_mode: str | None = None
+    host_weight_runtime_root: str | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/diffusion/data.py

@@ -100,6 +100,19 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
     return normalized
 
 
+def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
+    """Validate HWR policy without touching the configured storage domain.
+
+    Filesystem locality and store construction belong to the eligible loader
+    path. Keeping this check purely structural is what lets disabled and
+    AllGather DLO configurations retain zero HWR interaction.
+    """
+    if mode not in {"disabled", "preferred", "required"}:
+        raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required")
+    if mode != "disabled" and (not isinstance(root, str) or not root.strip()):
+        raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
+
+
 def parse_kv_cache_skip_selector(
     selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -766,6 +779,11 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Final-layout Host Weight Runtime policy. The loader only activates this
+    # for eligible no-AllGather DLO; all other configurations preserve their
+    # existing loader/storage path.
+    host_weight_runtime_mode: str = "disabled"
+    host_weight_runtime_root: str | None = None
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
@@ -1142,6 +1160,11 @@ class OmniDiffusionConfig:
         elif self.max_cpu_loras < 1:
             raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
 
+        validate_host_weight_runtime_options(
+            mode=self.host_weight_runtime_mode,
+            root=self.host_weight_runtime_root,
+        )
+
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(
                 "diffusers_load_kwargs and diffusers_call_kwargs are only "

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import contextlib
 import dataclasses
 import glob
@@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
+from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     build_checkpoint_mmap_plan,
@@ -125,7 +126,7 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ
     )
 
 
-class DiffusersPipelineLoader:
+class DiffusersPipelineLoader(HWRLoaderMixin):
     """Model loader that can load diffusers pipeline components from disk."""
 
     @dataclasses.dataclass
@@ -159,6 +160,9 @@ class DiffusersPipelineLoader:
         self.quant_config = od_config.quantization_config
         self.parallel_config = od_config.parallel_config
         self.host_weight_plan: HostWeightPlan | None = None
+        self._hwr_state: dict[str, object] | None = None
+        self._last_load_request: dict[str, object] | None = None
+        self._force_canonical_load = False
 
     def take_host_weight_plan(self) -> HostWeightPlan | None:
         """Transfer the loader-produced plan to the offload backend."""
@@ -469,6 +473,13 @@ class DiffusersPipelineLoader:
     ) -> nn.Module:
         """Load a model with the given configurations."""
         self.host_weight_plan = None
+        self._hwr_state = None
+        self._last_load_request = {
+            "load_device": load_device,
+            "load_format": load_format,
+            "custom_pipeline_name": custom_pipeline_name,
+            "device": device,
+        }
         if load_format is None:
             load_format = "default"
         # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
@@ -511,7 +522,34 @@ class DiffusersPipelineLoader:
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload:
+                hwr_state = None
+                if not self._force_canonical_load:
+                    try:
+                        hwr_state = self._resolve_hwr(
+                            model,
+                            ModuleDiscovery.discover(model),
+                            dist_offload=_dist_offload,
+                            use_allgather=_use_ag,
+                            load_format=load_format,
+                            sources=weight_sources,
+                        )
+                    except _HWRCommitError:
+                        from vllm_omni.host_weight_runtime import RuntimeMode
+
+                        mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled"))
+                        if mode is not RuntimeMode.PREFERRED:
+                            raise
+                        logger.warning(
+                            "HWR restore commit failed; discarding the model and retrying a fresh canonical load",
+                            exc_info=True,
+                        )
+                        del model
+                        return self.load_fresh_canonical_model()
+                self._hwr_state = hwr_state
+                hwr_active = hwr_state is not None
+                if hwr_active and hwr_state is not None:
+                    self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan"))
+                if _dist_offload and not hwr_active and not self._force_canonical_load:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -584,12 +622,43 @@ class DiffusersPipelineLoader:
                             self.load_weights(model)
                     self._process_weights_after_loading(model, target_device)
 
+                # A warm final-layout hit has already completed all
+                # byte-changing work through the restorer.  Shared runtime
+                # finalization happens once at the end for both cold and warm
+                # paths; the warm path never re-enters the ordinary
+                # materialization/finalization pipeline.
+
             if offload_after_quant:
                 model.to("cpu")
                 logger.info("Quantization complete, offloaded model back to CPU")
 
-        self._apply_skip_softmax_calibration(model)
-        return model.eval()
+        try:
+            self._apply_skip_softmax_calibration(model)
+            model = model.eval()
+            if self._hwr_state is not None:
+                warm_snapshot = self._hwr_state.get("warm_snapshot")
+                if warm_snapshot is not None:
+                    self._assert_final_layout_tensors_unchanged(model, cast(dict[str, tuple[int, str]], warm_snapshot))
+                self._publish_hwr_after_load(model, ModuleDiscovery.discover(model), self._hwr_state)
+        except Exception:
+            hwr_plan = self._hwr_state.get("plan") if self._hwr_state is not None else None
+            if isinstance(hwr_plan, HostWeightPlan):
+                carrier = hwr_plan.lease_carrier
+                if carrier is not None:
+                    carrier.close()
+                from vllm_omni.host_weight_runtime import RuntimeMode
+
+                mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled"))
+                if mode is RuntimeMode.PREFERRED:
+                    logger.warning(
+                        "HWR warm finalization failed; discarding the model and retrying a fresh canonical load",
+                        exc_info=True,
+                    )
+                    del model
+                    return self.load_fresh_canonical_model()
+            raise
+        self._attach_offload_startup_state(model)
+        return model
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:

vllm_omni/diffusion/model_loader/host_weight_loader.py

@@ -0,0 +1,479 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Host Weight Runtime integration for diffusion model loaders.
+
+This mixin keeps HWR identity, restore, publication, and startup-recovery
+policy out of the ordinary diffusers loader.  The host loader supplies the
+canonical source-preparation and weight-loading hooks; this module owns only
+the optional final-layout transaction.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import os
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, TensorBinding
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
+    from vllm_omni.diffusion.model_loader.host_weights.source_identity import PreparedWeightSource
+
+
+class _HWRCommitError(RuntimeError):
+    """A committed warm restore made the current model disposable."""
+
+
+class HWRLoaderMixin:
+    """Optional final-layout HWR behavior shared by diffusion loaders."""
+
+    @staticmethod
+    def _identity_value(value: object) -> object:
+        """Convert config objects into deterministic identity metadata."""
+        if value is None or isinstance(value, (bool, int, float, str)):
+            return value
+        if isinstance(value, torch.dtype):
+            return str(value)
+        if isinstance(value, dict):
+            return {str(key): HWRLoaderMixin._identity_value(item) for key, item in value.items()}
+        if isinstance(value, (list, tuple)):
+            return [HWRLoaderMixin._identity_value(item) for item in value]
+        to_dict = getattr(value, "to_dict", None)
+        if callable(to_dict):
+            try:
+                return HWRLoaderMixin._identity_value(to_dict())
+            except TypeError:
+                pass
+        if dataclasses.is_dataclass(value):
+            return HWRLoaderMixin._identity_value(dataclasses.asdict(value))
+        return f"{type(value).__module__}.{type(value).__qualname__}:{value!r}"
+
+    @staticmethod
+    def _identity_fingerprint(value: object) -> str:
+        from vllm_omni.host_weight_runtime import CanonicalJson
+
+        return hashlib.sha256(CanonicalJson.from_value(value).encoded).hexdigest()
+
+    @staticmethod
+    def _snapshot_final_layout_tensors(model: nn.Module, names: Iterable[str]) -> dict[str, tuple[int, str]]:
+        snapshot: dict[str, tuple[int, str]] = {}
+        for name in names:
+            parent_path, _, leaf_name = name.rpartition(".")
+            parent = model.get_submodule(parent_path)
+            tensor = parent._parameters.get(leaf_name)
+            if tensor is None:
+                tensor = parent._buffers.get(leaf_name)
+            if tensor is None or tensor.is_meta:
+                raise RuntimeError(f"cannot snapshot missing or meta final-layout tensor {name!r}")
+            value = tensor.detach()
+            if value.device.type != "cpu":
+                value = value.cpu()
+            value = value.contiguous()
+            digest = hashlib.sha256(memoryview(value.view(torch.uint8).numpy())).hexdigest()
+            snapshot[name] = (tensor.untyped_storage().data_ptr(), digest)
+        return snapshot
+
+    @classmethod
+    def _assert_final_layout_tensors_unchanged(
+        cls,
+        model: nn.Module,
+        snapshot: dict[str, tuple[int, str]],
+    ) -> None:
+        current = cls._snapshot_final_layout_tensors(model, snapshot)
+        changed = [name for name in snapshot if current[name] != snapshot[name]]
+        if changed:
+            raise RuntimeError(
+                "shared warm finalization changed restored final-layout tensor bytes or backing pointers: "
+                f"{changed[:5]}"
+            )
+
+    def _hwr_eligibility_mode(
+        self,
+        model: nn.Module,
+        modules: object,
+        *,
+        dist_offload: bool,
+        use_allgather: bool,
+        load_format: str,
+    ) -> object | None:
+        """Return the enabled HWR mode only after all zero-interaction gates."""
+        from vllm_omni.host_weight_runtime import RuntimeMode
+
+        raw_mode = getattr(self.od_config, "host_weight_runtime_mode", "disabled")
+        try:
+            mode = RuntimeMode(raw_mode)
+        except ValueError as exc:
+            raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") from exc
+
+        # These gates intentionally precede HWR imports, source preparation,
+        # identity construction, and store creation.
+        if mode is RuntimeMode.DISABLED or not dist_offload or use_allgather:
+            return None
+
+        parallel = self.parallel_config
+        reasons: list[str] = []
+        if load_format != "default":
+            reasons.append("load_format must be 'default'")
+        if bool(getattr(parallel, "use_hsdp", False)):
+            reasons.append("HSDP layouts are not supported by the final-layout BF16 consumer")
+        if self.quant_config is not None:
+            reasons.append("quantized layouts require a representation-specific HWR producer")
+        if getattr(self.od_config, "lora_path", None):
+            reasons.append("adapted weights are not reusable base-model artifacts")
+
+        dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
+        if not dit_modules:
+            reasons.append("no DiT modules were discovered")
+        elif any(
+            getattr(dit, "host_weight_restore_contract", None) is None
+            or not callable(getattr(dit, "validate_restored_host_weights", None))
+            for _, dit in dit_modules
+        ):
+            reasons.append("every owned DiT must declare the final-layout restore contract")
+
+        if reasons:
+            message = "; ".join(reasons)
+            if mode is RuntimeMode.REQUIRED:
+                raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
+            logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
+            return None
+        return mode
+
+    def _prepare_hwr_sources(
+        self,
+        model: nn.Module,
+        modules: object,
+        sources: Sequence[object],
+    ) -> tuple[PreparedWeightSource, ...]:
+        from vllm_omni.diffusion.model_loader.host_weights import (
+            ImplementationIdentity,
+            PreparedWeightSource,
+            WeightSourceKind,
+        )
+
+        dit_prefixes = tuple(f"{name}." for name in getattr(modules, "dit_names", ()))
+        selected_sources = tuple(
+            source
+            for source in sources
+            if not source.prefix or any(prefix.startswith(source.prefix) for prefix in dit_prefixes)
+        )
+        if not selected_sources:
+            raise ValueError("final-layout HWR requires canonical weight sources covering every DiT")
+
+        prepared: list[PreparedWeightSource] = []
+        for source in selected_sources:
+            resolved_root, weight_files, use_safetensors = self._prepare_weights(
+                source.model_or_path,
+                source.subfolder,
+                source.revision,
+                source.fall_back_to_pt,
+                source.allow_patterns_overrides,
+            )
+            adapter = self._get_checkpoint_adapter(model, source, use_safetensors)
+            adapter_identity = None
+            if adapter is not None:
+                adapter_name = f"{type(adapter).__module__}.{type(adapter).__qualname__}"
+                adapter_identity = ImplementationIdentity(
+                    implementation_id=adapter_name,
+                    version="1",
+                    fingerprint=self._identity_fingerprint({"adapter": adapter_name}),
+                )
+            source_kind = (
+                WeightSourceKind.LOCAL_PATH
+                if os.path.isdir(os.fspath(source.model_or_path))
+                else WeightSourceKind.HUGGING_FACE_HUB
+            )
+            prepared.append(
+                PreparedWeightSource(
+                    model_or_path=os.fspath(source.model_or_path),
+                    subfolder=source.subfolder,
+                    requested_revision=source.revision,
+                    prefix=source.prefix,
+                    resolved_root=Path(os.fspath(resolved_root)),
+                    weight_files=tuple(Path(os.fspath(path)) for path in weight_files),
+                    use_safetensors=use_safetensors,
+                    checkpoint_adapter=adapter_identity,
+                    source_kind=source_kind,
+                )
+            )
+        return tuple(prepared)
+
+    def _build_hwr_context(
+        self,
+        model: nn.Module,
+        modules: object,
+        *,
+        load_format: str,
+        sources: Sequence[object],
+    ) -> FinalLayoutIdentityContext:
+        from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+        from vllm_omni.diffusion.model_loader.host_weights import (
+            FINAL_LAYOUT_BF16_POLICY,
+            FinalLayoutLoaderIdentity,
+            FinalLayoutParallelIdentity,
+            FinalLayoutRequest,
+            ImplementationIdentity,
+            build_final_layout_identity,
+        )
+
+        parallel = self.parallel_config
+        tp_size = int(getattr(parallel, "tensor_parallel_size", 1))
+        try:
+            tp_rank = int(get_tensor_model_parallel_rank()) if tp_size > 1 else 0
+        except Exception:
+            tp_rank = int(getattr(parallel, "tensor_parallel_rank", 0))
+        sp_size = int(getattr(parallel, "sequence_parallel_size", 1) or 1)
+        loader_config = {
+            "dtype": str(self.od_config.dtype),
+            "model_class": f"{type(model).__module__}.{type(model).__qualname__}",
+            "model_config": self._identity_value(getattr(self.od_config, "tf_model_config", None)),
+            "load_format": load_format,
+            "quantization": self._identity_value(self.quant_config),
+        }
+        contracts = [
+            self._identity_value(getattr(dit, "host_weight_restore_contract")) for dit in getattr(modules, "dits", ())
+        ]
+        loader_identity = FinalLayoutLoaderIdentity(
+            implementation=ImplementationIdentity(
+                implementation_id="vllm-omni.diffusion.diffusers-loader",
+                version="final-layout-v1",
+                fingerprint=self._identity_fingerprint(
+                    {
+                        "loader": "diffusers-loader-final-layout-v1",
+                        "pipeline": f"{type(model).__module__}.{type(model).__qualname__}",
+                    }
+                ),
+            ),
+            model_config_fingerprint=self._identity_fingerprint(loader_config),
+            weight_transform_fingerprint=self._identity_fingerprint(
+                {
+                    "contracts": contracts,
+                    "transform": "diffusion-final-layout-loader-transforms-v1",
+                }
+            ),
+        )
+        semantic_parallel = FinalLayoutParallelIdentity(
+            tensor_parallel_size=tp_size,
+            tensor_parallel_rank=tp_rank,
+            sequence_parallel_size=sp_size,
+            ulysses_degree=int(getattr(parallel, "ulysses_degree", 1)),
+            ring_degree=int(getattr(parallel, "ring_degree", 1)),
+            allgather_degree=int(getattr(parallel, "allgather_degree", 1)),
+            ulysses_mode=str(getattr(parallel, "ulysses_mode", "strict")),
+            pipeline_parallel_size=int(getattr(parallel, "pipeline_parallel_size", 1)),
+            cfg_parallel_size=int(getattr(parallel, "cfg_parallel_size", 1)),
+            use_hsdp=bool(getattr(parallel, "use_hsdp", False)),
+            enable_expert_parallel=bool(getattr(parallel, "enable_expert_parallel", False)),
+        )
+        model_id = str(getattr(self.od_config, "model", "") or "")
+        if not model_id:
+            raise ValueError("final-layout HWR requires a canonical model identifier")
+        request = FinalLayoutRequest(
+            model_id=model_id,
+            loader=loader_identity,
+            parallel=semantic_parallel,
+            load_format=load_format,
+        )
+        prepared_sources = self._prepare_hwr_sources(model, modules, sources)
+        dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
+        return build_final_layout_identity(
+            model,
+            dit_modules=dit_modules,
+            prepared_sources=prepared_sources,
+            request=request,
+            policy=FINAL_LAYOUT_BF16_POLICY,
+        )
+
+    def _resolve_hwr(
+        self,
+        model: nn.Module,
+        modules: object,
+        *,
+        dist_offload: bool,
+        use_allgather: bool,
+        load_format: str,
+        sources: Sequence[object],
+    ) -> dict[str, object] | None:
+        """Resolve an eligible no-AllGather final-layout HWR transaction."""
+        from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutTensorRestorer
+        from vllm_omni.host_weight_runtime import (
+            HostWeightLeaseCarrier,
+            HostWeightRuntime,
+            HostWeightRuntimeConfig,
+            ProductionPolicy,
+            ResolutionOutcome,
+            RuntimeMode,
+            StorageDomainPolicy,
+        )
+        from vllm_omni.host_weight_runtime.filesystem import detect_storage_class
+
+        mode = self._hwr_eligibility_mode(
+            model,
+            modules,
+            dist_offload=dist_offload,
+            use_allgather=use_allgather,
+            load_format=load_format,
+        )
+        if mode is None:
+            return None
+        assert isinstance(mode, RuntimeMode)
+        expected_prefixes = frozenset(f"{name}." for name in getattr(modules, "dit_names", ()))
+        available_prefixes = frozenset(getattr(source, "prefix", "") for source in sources)
+        if not expected_prefixes <= available_prefixes:
+            message = "final-layout HWR requires one dedicated source prefix per owned DiT"
+            if mode is RuntimeMode.REQUIRED:
+                raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
+            logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
+            return None
+        overlapping_prefixes = frozenset(
+            prefix
+            for prefix in available_prefixes
+            if any(dit_prefix.startswith(prefix) for dit_prefix in expected_prefixes)
+        )
+        if not overlapping_prefixes <= expected_prefixes:
+            message = "final-layout HWR requires dedicated DiT sources and rejects mixed component sources"
+            if mode is RuntimeMode.REQUIRED:
+                raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
+            logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
+            return None
+        context = self._build_hwr_context(model, modules, load_format=load_format, sources=sources)
+        root_value = getattr(self.od_config, "host_weight_runtime_root", None)
+        if not isinstance(root_value, str) or not root_value.strip():
+            raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
+        root = Path(root_value).expanduser()
+        runtime = HostWeightRuntime.from_config(
+            HostWeightRuntimeConfig(
+                mode=mode,
+                domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)),
+                production=ProductionPolicy(
+                    allow_local_build=False,
+                    allow_post_load_publish=True,
+                ),
+            )
+        )
+        resolution = runtime.resolve(context.identity)
+        state: dict[str, object] = {
+            "mode": mode,
+            "context": context,
+            "runtime": runtime,
+            "outcome": resolution.report.outcome,
+        }
+        if resolution.report.outcome is ResolutionOutcome.FAILED:
+            failure = next(
+                (attempt.failure for attempt in reversed(resolution.report.attempts) if attempt.failure is not None),
+                None,
+            )
+            detail = failure.message if failure is not None else "resolution failed without a typed detail"
+            raise RuntimeError(f"Host Weight Runtime resolution failed: {detail}")
+        if resolution.report.outcome is not ResolutionOutcome.LOCAL_HIT:
+            return state
+
+        lease = resolution.lease
+        if lease is None:
+            raise RuntimeError("Host Weight Runtime returned LOCAL_HIT without a lease")
+        restorer = FinalLayoutTensorRestorer(context)
+        try:
+            restore_plan = restorer.plan_restore(model, lease)
+        except Exception:
+            lease.close()
+            if mode is RuntimeMode.REQUIRED:
+                raise
+            logger.warning("HWR warm restore planning failed; falling back to canonical loading", exc_info=True)
+            state["outcome"] = ResolutionOutcome.CANONICAL_FALLBACK
+            return state
+
+        try:
+            restore_plan.commit()
+        except Exception as exc:
+            lease.close()
+            raise _HWRCommitError(
+                "Host Weight Runtime restore commit failed; the partially restored model must be discarded"
+            ) from exc
+
+        try:
+            carrier = HostWeightLeaseCarrier(lease)
+            warm_snapshot = self._snapshot_final_layout_tensors(model, context.tensor_names)
+        except Exception as exc:
+            lease.close()
+            raise _HWRCommitError(
+                "Host Weight Runtime committed restore could not establish its startup ownership boundary"
+            ) from exc
+        source_metadata = context.identity.source.metadata.to_value()
+        target_bindings = source_metadata.get("target_bindings") if isinstance(source_metadata, dict) else None
+        planned_prefixes = frozenset(
+            binding["source_prefix"]
+            for binding in (target_bindings or ())
+            if isinstance(binding, dict) and isinstance(binding.get("source_prefix"), str)
+        )
+        if not planned_prefixes:
+            lease.close()
+            raise _HWRCommitError("Host Weight Runtime identity did not record owned canonical source prefixes")
+        state["plan"] = HostWeightPlan(
+            backing_kind="host_weight_runtime",
+            bindings={name: TensorBinding(name, "") for name in context.tensor_names},
+            planned_source_prefixes=planned_prefixes,
+            lease_carrier=carrier,
+        )
+        state["warm_snapshot"] = warm_snapshot
+        return state
+
+    def load_fresh_canonical_model(self) -> nn.Module:
+        """Reload a disposable startup model without HWR or checkpoint mmap."""
+        if self._last_load_request is None:
+            raise RuntimeError("cannot construct a fresh canonical model before the initial load")
+        request = dict(self._last_load_request)
+        self._force_canonical_load = True
+        try:
+            return self.load_model(**cast(dict[str, object], request))
+        finally:
+            self._force_canonical_load = False
+
+    def _publish_hwr_after_load(self, model: nn.Module, modules: object, state: dict[str, object] | None) -> None:
+        from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutBF16Producer
+        from vllm_omni.host_weight_runtime import PostLoadPublicationOutcome, ResolutionOutcome
+
+        if state is None or state.get("outcome") is not ResolutionOutcome.CANONICAL_FALLBACK:
+            return
+        context = state["context"]
+        runtime = state["runtime"]
+        dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
+        try:
+            producer = FinalLayoutBF16Producer(context, model, dit_modules)
+            report = runtime.publish_after_load(context.identity, producer=producer)
+        except Exception:
+            logger.warning("Host Weight Runtime post-load publication failed", exc_info=True)
+            return
+        if report.outcome is PostLoadPublicationOutcome.FAILED:
+            logger.warning("Host Weight Runtime post-load publication failed: %s", report.failure)
+
+    def _attach_offload_startup_state(self, model: nn.Module) -> None:
+        """Publish generic offload startup state on the loaded pipeline."""
+        if self.host_weight_plan is None:
+            return
+        from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
+
+        hwr_state = self._hwr_state
+        allow_retry = False
+        if hwr_state is not None and isinstance(hwr_state.get("plan"), HostWeightPlan):
+            from vllm_omni.host_weight_runtime import RuntimeMode
+
+            allow_retry = hwr_state.get("mode") is RuntimeMode.PREFERRED
+        attach_offload_startup_state(
+            model,
+            OffloadStartupState(
+                host_weight_plan=self.host_weight_plan,
+                fresh_model_loader=self.load_fresh_canonical_model if allow_retry else None,
+                allow_fresh_retry=allow_retry,
+            ),
+        )

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Loader-owned host-weight plans shared with diffusion offload backends."""
 
 from __future__ import annotations
@@ -9,6 +9,7 @@ import json
 import os
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import torch
 from safetensors import safe_open
@@ -21,6 +22,9 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
 
 logger = init_logger(__name__)
 
+if TYPE_CHECKING:
+    from vllm_omni.host_weight_runtime import HostWeightLeaseCarrier
+
 TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 
 
@@ -40,6 +44,7 @@ class HostWeightPlan:
     backing_kind: str
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
+    lease_carrier: HostWeightLeaseCarrier | None = None
 
 
 @dataclass(frozen=True)

vllm_omni/diffusion/offloader/__init__.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import torch
 from vllm.logger import init_logger
@@ -25,6 +25,7 @@ from .sequential_backend import (
     remove_sequential_offload,
     sequential_offload_component,
 )
+from .startup import OffloadStartupState, take_offload_startup_state
 from .tensor_utils import (
     dtype_size,
     is_dtensor,
@@ -52,6 +53,9 @@ __all__ = [
     "apply_distributed_block_hook",
     "remove_distributed_block_hook",
     "get_offload_backend",
+    "enable_offload_backend",
+    "OffloadStartupState",
+    "take_offload_startup_state",
     "get_offload_plan",
     "get_blocks_attr_names",
     "get_blocks_from_dit",
@@ -137,3 +141,55 @@ def get_offload_backend(
     else:
         logger.error("Unknown offload strategy: %s", config.strategy)
         return None
+
+
+def enable_offload_backend(
+    od_config: OmniDiffusionConfig,
+    pipeline: torch.nn.Module,
+    device: torch.device | None = None,
+) -> tuple[torch.nn.Module, OffloadBackend | None]:
+    """Create and enable the loader-selected backend transactionally.
+
+    The model runner only calls this generic offloader boundary. Loader-owned
+    host plans and optional fresh-model recovery callbacks stay inside the
+    startup state consumed here.
+    """
+
+    def enable_once(model: torch.nn.Module, state: OffloadStartupState | None) -> OffloadBackend | None:
+        backend: OffloadBackend | None = None
+        try:
+            backend = get_offload_backend(
+                od_config,
+                device=device,
+                host_weight_plan=state.host_weight_plan if state is not None else None,
+            )
+            if backend is not None:
+                logger.info("Enabling offloader backend: %s", backend.__class__.__name__)
+                backend.enable(model)
+            elif state is not None:
+                state.close_loader_ownership()
+            return backend
+        except Exception:
+            if backend is not None:
+                try:
+                    backend.disable()
+                except Exception:
+                    logger.exception("Failed to clean up the offload backend after startup failure")
+            if state is not None:
+                state.close_loader_ownership()
+            raise
+
+    startup_state = take_offload_startup_state(pipeline)
+    try:
+        return pipeline, enable_once(pipeline, startup_state)
+    except Exception:
+        if startup_state is None or not startup_state.allow_fresh_retry:
+            raise
+        assert startup_state.fresh_model_loader is not None
+        logger.warning(
+            "Loader-backed offload startup failed; retrying with a fresh canonical model",
+            exc_info=True,
+        )
+        del pipeline
+        pipeline = startup_state.fresh_model_loader()
+        return pipeline, enable_once(pipeline, take_offload_startup_state(pipeline))

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -30,6 +30,7 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
 )
+from vllm_omni.host_weight_runtime import HostWeightLease
 from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
@@ -902,6 +903,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = False
         self._using_rank_local_mmap = False
         self.host_weight_plan = host_weight_plan
+        self._host_weight_lease: HostWeightLease | None = None
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
 
     def load_resident_layers(self) -> None:
@@ -1405,6 +1407,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 buffer.data = buffer.data.to(self.device, non_blocking=True)
 
     def enable(self, pipeline: nn.Module) -> None:
+        """Enable DLO and make partial startup failures transactional."""
+        try:
+            self._enable(pipeline)
+        except BaseException:
+            try:
+                self.disable()
+            except BaseException:
+                logger.exception("DLO cleanup failed while handling an enable failure")
+            raise
+
+    def _enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("DistributedLayerwiseOffloadBackend already enabled")
             return
@@ -1443,16 +1456,33 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = self.host_weight_plan is not None
         self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
         if self._using_mmap:
-            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
+            if self.host_weight_plan.backing_kind == "host_weight_runtime":
+                carrier = self.host_weight_plan.lease_carrier
+                if carrier is None:
+                    raise RuntimeError("DLO received a Host Weight Runtime plan without a lease carrier")
+                self._host_weight_lease = carrier.take()
+                if self._host_weight_lease.closed:
+                    raise RuntimeError("DLO received a closed Host Weight Runtime lease")
+                # The final-layout restorer has already rebound the model to
+                # immutable host tensors.  Treat those tensors as mmap-like
+                # sources for the bounded two-slot rank-local staging path;
+                # do not reinitialize or rematerialize the restored model.
+                self._using_rank_local_mmap = True
+                logger.info(
+                    "DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging",
+                    self._host_weight_lease.provenance.resolution_id,
+                )
+            elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
+                self._load_weights_via_mmap(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+            else:
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
-            self._load_weights_via_mmap(
-                pipeline,
-                modules,
-                self.host_weight_plan,
-            )
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: checkpoint pages are "
+                    "DLO rank-local host storage enabled: source pages are "
                     "node-shared; each worker owns only two bounded host staging slots"
                 )
         else:
@@ -1722,12 +1752,19 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._cleanup_after_loading()
 
     def _release_mmap_handles(self) -> None:
-        """Release safetensors mmap file handles."""
+        """Release source handles and the transport-owned HWR lease."""
         self._mmap_transforms_by_tensor_id.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
             logger.info("Released safetensors mmap file handles")
+        lease = self._host_weight_lease
+        self._host_weight_lease = None
+        if lease is not None and not lease.closed:
+            lease.close()
+            logger.info("Released Host Weight Runtime lease %s", lease.provenance.resolution_id)
+        if self.host_weight_plan is not None and self.host_weight_plan.lease_carrier is not None:
+            self.host_weight_plan.lease_carrier.close()
 
     def _cleanup_after_loading(self) -> None:
         """Synchronize and release freed device/CPU memory after sharding."""
@@ -1743,7 +1780,25 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             pass
 
     def disable(self) -> None:
-        if not self.enabled and not hasattr(self, "_mmap_file_cache"):
+        has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
+        has_carrier = (
+            self.host_weight_plan is not None
+            and self.host_weight_plan.lease_carrier is not None
+            and not self.host_weight_plan.lease_carrier.closed
+        )
+        has_partial_hooks = bool(
+            self._blocks
+            or self._all_hook_groups
+            or getattr(self, "_on_demand_handles", ())
+            or self._resident_layer_group is not None
+        )
+        if (
+            not self.enabled
+            and not hasattr(self, "_mmap_file_cache")
+            and not has_open_lease
+            and not has_carrier
+            and not has_partial_hooks
+        ):
             return
 
         if self._using_rank_local_mmap:

vllm_omni/diffusion/offloader/startup.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Generic startup handoff between model loaders and offload backends."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
+
+_STARTUP_STATE_ATTR = "_omni_offload_startup_state"
+
+
+@dataclass
+class OffloadStartupState:
+    """Loader-owned state consumed by the offloader startup boundary."""
+
+    host_weight_plan: HostWeightPlan | None = None
+    fresh_model_loader: Callable[[], nn.Module] | None = None
+    allow_fresh_retry: bool = False
+
+    def __post_init__(self) -> None:
+        if self.allow_fresh_retry != (self.fresh_model_loader is not None):
+            raise ValueError("fresh-model retry policy requires exactly one retry callback")
+
+    def close_loader_ownership(self) -> None:
+        """Release a plan that never reached a backend."""
+        if self.host_weight_plan is None:
+            return
+        carrier = self.host_weight_plan.lease_carrier
+        if carrier is not None:
+            carrier.close()
+
+
+def attach_offload_startup_state(model: nn.Module, state: OffloadStartupState) -> None:
+    """Attach one process-local startup handoff to a loaded pipeline."""
+    previous = getattr(model, _STARTUP_STATE_ATTR, None)
+    if previous is not None:
+        previous.close_loader_ownership()
+    setattr(model, _STARTUP_STATE_ATTR, state)
+
+
+def take_offload_startup_state(model: nn.Module) -> OffloadStartupState | None:
+    """Take and remove the loader handoff from a pipeline exactly once."""
+    state = getattr(model, _STARTUP_STATE_ATTR, None)
+    if state is None:
+        return None
+    delattr(model, _STARTUP_STATE_ATTR)
+    return state
+
+
+__all__ = [
+    "OffloadStartupState",
+    "attach_offload_startup_state",
+    "take_offload_startup_state",
+]

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -49,7 +49,7 @@ from vllm_omni.diffusion.models.interface import (
     supports_prompt_update,
     supports_step_execution,
 )
-from vllm_omni.diffusion.offloader import get_offload_backend
+from vllm_omni.diffusion.offloader import enable_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
@@ -336,15 +336,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 f"{self.od_config.model_class_name} does not support that contract."
             )
 
-        # Apply CPU offloading
-        self.offload_backend = get_offload_backend(
+        # The offloader owns loader-plan handoff and startup recovery. The
+        # runner only receives the pipeline/backend pair that is ready to use.
+        self.pipeline, self.offload_backend = enable_offload_backend(
             self.od_config,
+            self.pipeline,
             device=self.device,
-            host_weight_plan=model_loader.take_host_weight_plan(),
         )
-        if self.offload_backend is not None:
-            logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}")
-            self.offload_backend.enable(self.pipeline)
 
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:

vllm_omni/engine/arg_utils.py

@@ -559,6 +559,8 @@ class OrchestratorArgs:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0
+    host_weight_runtime_mode: str = "disabled"
+    host_weight_runtime_root: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1059,6 +1059,8 @@ class AsyncOmniEngine:
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
+            "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
+            "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
                 "regional"

vllm_omni/entrypoints/cli/serve.py

@@ -738,6 +738,28 @@ class OmniServeCommand(CLISubcommand):
             help="Keep this many leading main-DiT blocks resident on the device "
             "while distributed layerwise offload streams the remaining blocks.",
         )
+        omni_config_group.add_argument(
+            "--host-weight-runtime-mode",
+            choices=("disabled", "preferred", "required"),
+            default="disabled",
+            help=(
+                "Host Weight Runtime policy for eligible no-AllGather DLO: "
+                "disabled does not consult HWR; preferred restores an exact hit "
+                "or canonically loads and publishes on a miss; required restores "
+                "an exact hit or fails startup. Populate a required store with "
+                "preferred first."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--host-weight-runtime-root",
+            type=str,
+            default=None,
+            help=(
+                "Writable node-local Host Weight Runtime store shared by workers "
+                "in one storage domain. Required for preferred and required; use "
+                "the same persistent path for population and serving."
+            ),
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/host_weight_runtime/__init__.py

@@ -35,7 +35,7 @@ from .inspection import (
     DomainInspection,
     HostWeightInspector,
 )
-from .lease import HostWeightLease, LeaseProvenance, MappedHostRegion
+from .lease import HostWeightLease, HostWeightLeaseCarrier, LeaseProvenance, MappedHostRegion
 from .manifest import (
     FileManifestEntry,
     ManifestValidationError,
@@ -92,6 +92,7 @@ __all__ = [
     "HostWeightError",
     "HostWeightFailure",
     "HostWeightLease",
+    "HostWeightLeaseCarrier",
     "HostWeightResolution",
     "HostWeightInspector",
     "HostWeightRuntime",

vllm_omni/host_weight_runtime/filesystem/store.py

@@ -863,7 +863,11 @@ class FilesystemHostWeightStore:
                 manifest, _ = writer.finalize(identity, metadata)
                 artifact_bytes = self._tree_bytes(staging)
                 self._check_capacity(artifact_bytes, 0)
-                os.chmod(staging, 0o555)
+                # Keep the staging directory writable until the atomic rename.
+                # overlayfs rejects renaming a non-writable source directory;
+                # the artifact lock keeps readers out while the published
+                # directory is hardened below.
+                os.chmod(staging, 0o755)
                 _fsync_directory(staging)
             except HostWeightError as exc:
                 if writer is not None:
@@ -934,6 +938,11 @@ class FilesystemHostWeightStore:
                         self._remove_tree(staging)
                     else:
                         os.replace(staging, entry_path)
+                        try:
+                            os.chmod(entry_path, 0o555)
+                        except OSError:
+                            self._quarantine_locked(identity.key, reason="publication_hardening_failure")
+                            raise
                         _fsync_directory(self.artifacts_dir)
                         _fsync_directory(self.tmp_dir)
                     self._remove_deny_locked(identity.key)

vllm_omni/host_weight_runtime/lease.py

@@ -140,4 +140,60 @@ class HostWeightLease:
             self.close()
 
 
-__all__ = ["HostWeightLease", "LeaseProvenance", "MappedHostRegion"]
+class HostWeightLeaseCarrier:
+    """Process-local, single-take ownership carrier for a host-weight lease.
+
+    The loader creates the carrier after a restore transaction has committed.
+    The runner/backend takes the lease exactly once before starting any
+    asynchronous transport work.  If setup fails before that handoff, closing
+    the carrier releases the loader-owned lease.
+    """
+
+    def __init__(self, lease: HostWeightLease) -> None:
+        if not isinstance(lease, HostWeightLease):
+            raise TypeError("HostWeightLeaseCarrier requires one HostWeightLease")
+        if lease.closed:
+            raise ValueError("cannot carry a closed HostWeightLease")
+        self._lease: HostWeightLease | None = lease
+        self._taken = False
+        self._lock = threading.RLock()
+
+    @property
+    def taken(self) -> bool:
+        with self._lock:
+            return self._taken
+
+    @property
+    def closed(self) -> bool:
+        with self._lock:
+            return self._lease is None or self._lease.closed
+
+    def take(self) -> HostWeightLease:
+        """Transfer the lease to the transport owner exactly once."""
+        with self._lock:
+            if self._taken:
+                raise RuntimeError("HostWeightLeaseCarrier was already taken")
+            lease = self._lease
+            if lease is None or lease.closed:
+                raise RuntimeError("HostWeightLeaseCarrier does not own an open lease")
+            self._lease = None
+            self._taken = True
+            return lease
+
+    def close(self) -> None:
+        """Close the lease if ownership has not been transferred."""
+        with self._lock:
+            lease = self._lease
+            self._lease = None
+        if lease is not None:
+            lease.close()
+
+    def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn:
+        raise TypeError("HostWeightLeaseCarrier is process-local and cannot be serialized")
+
+    def __del__(self) -> None:
+        with contextlib.suppress(Exception):
+            self.close()
+
+
+__all__ = ["HostWeightLease", "HostWeightLeaseCarrier", "LeaseProvenance", "MappedHostRegion"]