[Diffusion] Register HWR mmap for direct DLO H2D (#6591)

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

@@ -196,16 +196,15 @@ adopt it.
 
 ### 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.
+> **Status:** PR2 landed in
+> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
+> registered direct H2D plus generic HWR publication/source-digest
+> optimizations without changing AllGather behavior.
 
 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.
+`--host-weight-runtime-root <node-local-root>`.
 
 The modes express operator fallback policy, not different artifact formats:
 
@@ -254,6 +253,44 @@ outcomes are authoritative:
 Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
 not selected.
 
+Local source identity and publication avoid repeated full-payload passes while
+preserving the same semantic contract:
+
+- Canonical files without a trusted immutable Hub blob identity are hashed in
+  parallel. Their content IDs are cached per storage domain under the HWR root
+  and reused only when the path, inode, size, timestamps, symlink target, and
+  cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
+  cache-I/O failure falls back to hashing the canonical file directly.
+- Producers that emit payload bytes in canonical storage-key order may declare
+  `ordered=True`. The filesystem writer then computes file and tensor SHA256
+  values during the write and overlaps payload `fsync` with later producer
+  work. Producers without that guarantee retain unordered writes and use a
+  parallel readback checksum fallback.
+- Manifest and `READY` publication still wait for every payload `fsync` and
+  checksum to complete. These optimizations do not weaken artifact validation
+  or change lease ownership.
+
+#### Registered mmap transport
+
+After DLO takes an HWR lease, the no-AllGather backend may register the lease's
+complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
+Registration is transport state: it does not change artifact identity, the
+store, tensor ownership, or H2D payload. On success, each tensor view copies
+directly into the existing rotating HBM block buffers and the two private host
+staging slots are not allocated.
+
+`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
+over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
+memory policy, unsupported platform/capability, over-budget mapping, or fully
+rolled-back registration error selects the existing two-slot staging path. A
+partial registration that cannot be rolled back aborts startup because closing
+the lease would unmap memory still owned by the platform.
+
+Direct checkpoint mmap remains unchanged and continues to use staging. It may
+require loader-owned per-block transforms, while the HWR artifact already
+contains final runtime bytes. DLO AllGather never receives an HWR final-layout
+lease and therefore never enters this registration path.
+
 #### Pre-service transaction boundary
 
 The startup transaction does not end at restore commit. A lease-backed model is
@@ -317,6 +354,7 @@ loader owns
   -> backend takes once before asynchronous work
   -> backend drains pending H2D work
   -> backend releases hooks, staging, and model references
+  -> backend unregisters every mapped range
   -> backend closes lease
 ```
 
@@ -326,13 +364,15 @@ 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.
+initial prefetch. Successful registration copies directly from immutable HWR
+views; otherwise the backend uses the existing two bounded rank-local staging
+slots. Transactional cleanup drains device work, removes source references,
+unregisters mappings, and only then closes the lease. An unregistration failure
+retains both registration and lease for retry/process teardown. Preferred mode
+then uses the runner's fresh canonical retry; required mode propagates the
+failure.
 
-#### PR2 promotion gates
+#### PR2 and PR3 promotion gates
 
 - Warm hit performs zero ordinary DiT materialization and zero producer calls.
 - Shared warm finalization changes neither restored bytes nor backing pointers.
@@ -344,6 +384,10 @@ propagates the failure.
 - 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.
+- Registration is all-or-nothing; rollback/unregistration failure never closes
+  a mapping still owned by the platform.
+- Registered HWR transport bypasses host staging while preserving output and
+  H2D payload; unsupported registration retains the bounded staging path.
 - Compatible checkpoint mmap remains an unchanged benchmark and control path.
 - Prewarm uses one matching producer cohort per TP/SP topology and storage
   domain.
@@ -394,11 +438,12 @@ 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.
 
-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.
+The Host Weight Runtime representation, publication, and no-AllGather consumer
+contracts are merged; see
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
+[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
+direct H2D is an optional transport layer over that merged lease contract.
 
 ## Validation coverage
 
@@ -411,8 +456,12 @@ Current source-level validation includes:
 - ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
   sharding of finalized weights and scales;
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
+- ordered publication hashing, overlapped durability, unordered parallel
+  checksum fallback, and node-local source-digest reuse/invalidation;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;
+- read-only registration preflight, direct source-to-device copies, safe
+  fallback, partial rollback, retryable unregistration, and lease ordering;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - bounded component allocator-cache retention, conservative/forced release,

docs/design/module/host_weight_runtime.md

@@ -16,7 +16,7 @@ validation_paths:
   - tests/host_weight_runtime/**
 upstream_refs:
   - https://github.com/vllm-project/vllm-omni/issues/6414
-last_reviewed: 2026-08-21
+last_reviewed: 2026-08-25
 ---
 
 # Host Weight Runtime
@@ -84,6 +84,13 @@ another. `WeightArtifactIdentity` includes:
 - producer version, implementation fingerprint, manifest schema, and restorer
   schema.
 
+Loader integrations may persist content digests for canonical source shards in
+the selected storage domain. Reuse requires the same path, inode, size,
+timestamps, symlink target, and cache-record checksum. Corrupt records are
+rebuilt, and coordination or cache-I/O failure falls back to hashing the
+canonical source directly. This is an identity-computation optimization, not a
+second source of truth or an artifact substitution rule.
+
 The identity excludes transfer policy such as registered mmap versus private
 pinned staging. Those paths move the same host representation and must not
 produce duplicate artifacts.
@@ -237,6 +244,12 @@ Every publication computes full file and tensor hashes. Local lookup supports:
   and
 - `full_checksum`, which additionally reads and hashes every payload file.
 
+An ordered producer may compute file and tensor digests while payload bytes are
+written and overlap each closed payload's `fsync` with later producer work. An
+unordered producer retains parallel readback hashing. In both cases manifest
+and `READY` publication wait for every payload digest and `fsync`; the atomic
+publication and validation contract is unchanged.
+
 Filesystem verity is modeled but unsupported. Invalid artifacts receive an
 external deny marker before their shared lookup lock is released; a build owner
 then quarantines the immutable directory under the exclusive lock.

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -65,6 +65,7 @@ omni = Omni(
 | `--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 |
+| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
 
 ## Host-weight loading
 
@@ -178,12 +179,48 @@ 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.
+DiT final tensors without ordinary DiT materialization. DLO then attempts to
+register the immutable mapping for direct H2D; if registration is unavailable,
+it streams 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.
+
+For local canonical checkpoints, the first eligible worker may hash source
+shards to establish immutable identity. HWR caches those digests in the same
+node-local storage domain and validates file metadata before reuse, so later
+workers normally avoid repeating that read. Cold BF16 publication also hashes
+ordered payloads as they are written and overlaps payload durability work with
+later shards; this changes startup work only, not artifact contents or runtime
+H2D behavior.
+
+#### Registered direct H2D
+
+Registration is attempted automatically only for an eligible warm HWR hit when
+`pin_cpu_memory` is enabled. A successful path is:
+
+```text
+shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
+```
+
+It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
+the two host staging slots. It does not reduce the H2D payload or make GPU
+kernels access host memory directly.
+
+CUDA requires read-only host-registration capability for the immutable HWR
+mapping. Unsupported capability, a positive registration limit smaller than
+the complete page-aligned mapping, or a safely rolled-back registration error
+falls back to bounded staging. Programmatic configurations can set
+`pin_cpu_memory=False` to disable registration explicitly. A successful
+registration locks the complete mapped range in host memory for that worker's
+lifetime, so use
+`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
+required.
+
+Shutdown drains H2D work, releases hook/source references, unregisters every
+range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
+retain their existing transfer paths.
 
 When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
 not resolve HWR sources or construct its store.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlanResult,
     TensorBinding,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 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
@@ -173,6 +174,15 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
         str(canonical_root / "model.safetensors"),
     )
     store_root = tmp_path / "hwr-store"
+    hash_calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal hash_calls
+        hash_calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
 
     def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
         loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
@@ -205,6 +215,8 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     assert warm_plan is not None
     assert warm_plan.lease_carrier is not None
     warm_plan.lease_carrier.close()
+    assert hash_calls == 1
+    assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
 
 
 def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -6,6 +6,8 @@ from __future__ import annotations
 
 import dataclasses
 import gc
+import json
+import threading
 from collections.abc import Sequence
 from pathlib import Path
 from types import SimpleNamespace
@@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import (
     FinalLayoutRequest,
     FinalLayoutTensorRestorer,
     ImplementationIdentity,
+    NodeSourceDigestCache,
     PreparedWeightSource,
     WeightSourceKind,
     build_final_layout_identity,
 )
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.model_loader.host_weights.contracts import (
     FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
     FinalLayoutArtifactSpec,
@@ -380,6 +384,7 @@ def _identity(
     *,
     request: FinalLayoutRequest | None = None,
     policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
+    source_digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutIdentityContext:
     transformer = model.get_submodule("transformer")
     return build_final_layout_identity(
@@ -388,6 +393,7 @@ def _identity(
         prepared_sources=(source,),
         request=request or _request(),
         policy=policy,
+        source_digest_cache=source_digest_cache,
     )
 
 
@@ -564,6 +570,131 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path
     assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
 
 
+def test_node_source_digest_cache_reuses_recovers_and_invalidates(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache_root = tmp_path / "store"
+    calls = 0
+    original_sha256 = source_identity_module._sha256_file
+
+    def counted_sha256(path: Path, state: object) -> str:
+        nonlocal calls
+        calls += 1
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
+    first = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    second = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+
+    assert second.identity == first.identity
+    assert calls == 1
+
+    cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
+    cache_entry.write_text("not-json", encoding="utf-8")
+    recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert recovered.identity == first.identity
+    assert calls == 2
+
+    cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
+    cache_document["content_id"] = f"sha256:{'0' * 64}"
+    cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
+    checksum_recovered = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert checksum_recovered.identity == first.identity
+    assert calls == 3
+
+    source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
+    changed = _identity(
+        model,
+        source,
+        source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+    )
+    assert changed.identity != first.identity
+    assert calls == 4
+
+
+def test_source_identity_hashes_checkpoint_shards_in_parallel(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    second_file = source.resolved_root / "model-00002.safetensors"
+    second_file.write_bytes(b"second-checkpoint-shard")
+    source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
+    barrier = threading.Barrier(2)
+    thread_ids: set[int] = set()
+    thread_ids_lock = threading.Lock()
+    original_sha256 = source_identity_module._sha256_file
+
+    def synchronized_sha256(path: Path, state: object) -> str:
+        with thread_ids_lock:
+            thread_ids.add(threading.get_ident())
+        barrier.wait(timeout=5)
+        return original_sha256(path, state)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
+    _identity(model, source)
+
+    assert len(thread_ids) == 2
+
+
+def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+
+    def unavailable_lock(*_args: object, **_kwargs: object) -> object:
+        raise source_identity_module.FileLockTimeoutError("busy")
+
+    monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+    assert not tuple(cache.entries.glob("*.json"))
+
+
+def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = _TinyPipeline()
+    source = _prepared_source(tmp_path)
+    original_mkdir = Path.mkdir
+
+    def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
+        if "source-digests-v1" in path.parts:
+            raise OSError("injected cache directory failure")
+        original_mkdir(path, *args, **kwargs)  # type: ignore[arg-type]
+
+    monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
+    cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+    context = _identity(model, source, source_digest_cache=cache)
+
+    assert context.identity.source.revision.startswith("content-")
+
+
 def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
     model = _TinyPipeline()
     source_root = tmp_path / "local-symlink-source"

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import ctypes
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    _coalesce_ranges,
+    _mapped_regions,
+)
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationBudgetError,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+    register_host_mappings,
+)
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+        *,
+        read_only_supported: bool = True,
+        attribute_result: int = 0,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self._read_only_supported = read_only_supported
+        self._attribute_result = attribute_result
+        self.attribute_queries: list[tuple[int, int]] = []
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
+        self.attribute_queries.append((attribute, device))
+        ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported)
+        return self._attribute_result
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def _region(
+    file_name: str,
+    address: int,
+    length: int,
+    *,
+    read_only: bool = True,
+) -> MappedHostRegion:
+    return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only)
+
+
+@pytest.fixture(autouse=True)
+def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
+
+
+def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    regions = (_region("weights", 0x1000, 4096),)
+
+    def create(actual_regions, *, max_bytes):
+        assert actual_regions is regions
+        assert max_bytes == 4096
+        return sentinel
+
+    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+    assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel
+    with pytest.raises(HostRegistrationError, match="not supported on cpu"):
+        register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None)
+
+
+def test_ranges_align_per_file_without_merging_adjacent_files() -> None:
+    assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+    )
+    assert _mapped_regions(
+        (
+            _region("first.safetensors", 0x1000, 4096),
+            _region("second.safetensors", 0x2000, 4096),
+        )
+    ) == (
+        registration_module._AddressRange(0x1000, 0x2000),
+        registration_module._AddressRange(0x2000, 0x3000),
+    )
+
+
+def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(HostRegistrationError, match="is writable"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None)
+    with pytest.raises(HostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096)
+    assert runtime.registered == []
+
+
+def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], read_only_supported=False)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(HostRegistrationError, match="does not support read-only"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+    assert runtime.attribute_queries == [
+        (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
+    ]
+    assert runtime.registered == []
+
+
+def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], attribute_result=7)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    consumed: list[tuple[object, int]] = []
+    monkeypatch.setattr(
+        registration_module,
+        "_consume_last_cuda_error",
+        lambda actual_runtime, error: consumed.append((actual_runtime, error)),
+    )
+
+    with pytest.raises(HostRegistrationError, match="error-7"):
+        CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+    assert consumed == [(runtime, 7)]
+
+
+def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+    with pytest.raises(HostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+    assert runtime.unregistered == [0x1000]
+
+
+def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0, 7], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+    with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error:
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+    registration = error.value.active_registration
+    assert registration is not None
+    assert registration.close() == ()
+    assert runtime.unregistered == [0x1000, 0x1000]
+
+
+def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None:
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+    registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None)
+
+    assert registration.total_bytes == 4096
+    assert registration.region_count == 1
+    assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",)
+    assert registration.close() == ()
+    assert runtime.unregistered == [0x1000, 0x1000]
+
+
+@pytest.mark.parametrize("pending_error", [0, 801])
+def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
+    monkeypatch: pytest.MonkeyPatch,
+    pending_error: int,
+) -> None:
+    runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error)
+    monkeypatch.setattr(
+        registration_module.ctypes,
+        "CDLL",
+        lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"),
+    )
+
+    registration_module._consume_last_cuda_error(runtime, 801)
+
+
+def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None:
+    global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801)
+    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime)
+
+    registration_module._consume_last_cuda_error(SimpleNamespace(), 801)
+
+
+def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    def missing_global_symbol(_name: object) -> None:
+        raise OSError("cudart symbols are local")
+
+    monkeypatch.setattr(
+        registration_module.ctypes,
+        "CDLL",
+        missing_global_symbol,
+    )
+
+    with pytest.raises(HostRegistrationError, match="cannot clear") as error:
+        CudaHostRegistration.create(
+            (
+                _region("first", 0x1000, 4096),
+                _region("second", 0x9000, 4096),
+            ),
+            max_bytes=None,
+        )
+
+    assert not isinstance(error.value, HostRegistrationCleanupError)
+    assert runtime.unregistered == [0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -5,6 +5,7 @@
 
 import gc
 import json
+import weakref
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -17,6 +18,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
 from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     TensorBinding,
@@ -33,12 +35,14 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadHook,
     PinnedResidentLayerGroup,
 )
+from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 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.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -703,11 +707,15 @@ class _HWRPipeline(nn.Module):
 
 
 class _FakeHostWeightLease:
-    def __init__(self, resolution_id: str):
+    def __init__(self, resolution_id: str, events: list[str] | None = None):
         self.closed = False
         self.provenance = SimpleNamespace(resolution_id=resolution_id)
+        self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
+        self._events = events
 
     def close(self):
+        if self._events is not None:
+            self._events.append("lease")
         self.closed = True
 
 
@@ -757,6 +765,203 @@ def _hwr_backend():
     return _HWRPipeline(), backend, carrier, lease
 
 
+def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
+    current_block = nn.Linear(2, 2, bias=False)
+    next_block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    next_block.weight.data.copy_(expected)
+    hook = DistributedLayerwiseOffloadHook(
+        next_block=next_block,
+        device=torch.device("cpu"),
+        dp_group=None,
+        dp_size=1,
+        rank=0,
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    hook.initialize_hook(current_block)
+    hook.registered_mmap = True
+    hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging")  # type: ignore[method-assign]
+
+    hook.prefetch_layer(slot=0, non_blocking=True)
+
+    assert torch.equal(next_block.weight, expected)
+
+
+def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
+    block = nn.Linear(2, 2, bias=False)
+    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+    block.weight.data.copy_(expected)
+    group = PinnedResidentLayerGroup(
+        [block],
+        device=torch.device("cpu"),
+        copy_stream=DummyStream(),
+        pin_memory=False,
+        rank_local_mmap=True,
+    )
+    group.registered_mmap = True
+    group._cpu_staging_buffers.clear()
+
+    group.load()
+
+    assert torch.equal(block.weight, expected)
+
+
+def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    plan, _, lease = _fake_hwr_plan()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+            dlo_host_registration_limit_gib=1.5,
+        ),
+        torch.device("cuda"),
+        host_weight_plan=plan,
+    )
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    calls: list[tuple[object, int | None]] = []
+
+    class Registration:
+        total_bytes = 4096
+        region_count = 1
+
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            return ()
+
+    def register(regions, *, device, max_bytes):
+        assert device == torch.device("cuda")
+        calls.append((regions, max_bytes))
+        return Registration()
+
+    monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+    pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
+
+    assert backend._try_register_hwr_mmap((pinned_source,))  # type: ignore[arg-type]
+    assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend._release_registered_mmap()
+
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+
+
+def test_hwr_registration_budget_validation_is_transport_scoped():
+    assert (
+        validate_dlo_host_registration_options(
+            limit_gib=1.5,
+            enable_dlo=True,
+            use_allgather=False,
+            hwr_mode="preferred",
+        )
+        == 1.5
+    )
+    invalid = (
+        dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
+        dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
+    )
+    for options in invalid:
+        with pytest.raises(ValueError):
+            validate_dlo_host_registration_options(**options)
+
+
+def test_unregistration_precedes_lease_close_and_retries_failures(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    events: list[str] = []
+    responses = iter([("busy",), ()])
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return next(responses)
+
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retry", events)
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = Registration()
+    backend._using_rank_local_mmap = True
+    backend._using_registered_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    assert not lease.closed
+    assert backend._host_registration is not None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+    backend.disable()
+
+    assert lease.closed
+    assert backend._host_registration is None
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+    assert events == ["unregister", "unregister", "lease"]
+
+
+def test_failed_unregistration_retains_lease_after_backend_is_dropped(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
+    events: list[str] = []
+
+    class Registration:
+        @staticmethod
+        def close() -> tuple[str, ...]:
+            events.append("unregister")
+            return ("busy",)
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=True,
+            dp_size=1,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    lease = _FakeHostWeightLease("retained", events)
+    registration = Registration()
+    backend._host_weight_lease = lease  # type: ignore[assignment]
+    backend._host_registration = registration
+    backend._using_rank_local_mmap = True
+    backend.enabled = True
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    lease_ref = weakref.ref(lease)
+
+    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+        backend.disable()
+    del backend
+    del lease
+    gc.collect()
+
+    retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
+    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
+    assert lease_ref() is retained_lease
+    assert not retained_lease.closed
+
+    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+    retained_lease.close()
+
+
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 

tests/diffusion/test_worker_wrapper_base.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
 
 """
 Unit tests for WorkerWrapperBase class.
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+        events: list[str] = []
+        offload_backend = mocker.Mock()
+        offload_backend.disable.side_effect = lambda: events.append("offload")
+        kv_manager = mocker.Mock()
+        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+        destroy = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
+            side_effect=lambda: events.append("distributed"),
+        )
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_manager,
+        )
+
+        worker.shutdown()
+
+        assert events == ["offload", "kv", "distributed"]
+        destroy.assert_called_once_with()
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -378,6 +378,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
             "preferred",
             "--host-weight-runtime-root",
             "/var/cache/vllm-omni/hwr",
+            "--dlo-host-registration-limit-gib",
+            "80",
         ]
     )
 
@@ -387,8 +389,10 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
 
     assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
     assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
     assert engine_args["host_weight_runtime_mode"] == "preferred"
     assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+    assert engine_args["dlo_host_registration_limit_gib"] == 80
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

tests/host_weight_runtime/test_filesystem_store.py

@@ -57,7 +57,9 @@ from vllm_omni.host_weight_runtime.filesystem import (
     detect_storage_class,
 )
 from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
+from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module
 from vllm_omni.host_weight_runtime.filesystem.locks import FileLock
+from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -110,6 +112,150 @@ def _open_fd_targets() -> tuple[str, ...]:
     return tuple(targets)
 
 
+def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter:
+    staging = tmp_path / "staging"
+    staging.mkdir()
+    return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None)
+
+
+def test_ordered_writer_hashes_payload_during_writes(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    weight = _weight()
+    scale = torch.tensor([1.0, 2.0])
+
+    def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object:
+        raise AssertionError("ordered payload must not be rescanned")
+
+    monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan)
+    with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+        # The writer's ordered contract follows canonical storage-key order.
+        output.write_tensor("layer.scale", scale)
+        output.write_rows("layer.weight", 0, weight[:2])
+        output.write_rows("layer.weight", 2, weight[2:])
+
+    manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    payload = writer.staging_dir / "weights.safetensors"
+    assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest()
+    expected_tensor_digests = {
+        "layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(),
+        "layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(),
+    }
+    assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests
+
+
+def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None:
+    writer = _direct_writer(tmp_path)
+    output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+
+    with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"):
+        output.write_tensor("layer.weight", _weight())
+
+    output.abort()
+    assert not list(writer.staging_dir.iterdir())
+
+
+def test_payload_fsync_overlaps_producer_work(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    started = threading.Event()
+    release = threading.Event()
+    closed = threading.Event()
+    original_fsync = filesystem_writer_module._fsync_payload
+
+    def blocking_fsync(fd: int) -> None:
+        started.set()
+        if not release.wait(timeout=5):
+            raise TimeoutError("test did not release payload fsync")
+        original_fsync(fd)
+
+    monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync)
+    output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+    output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+    output.write_tensor("layer.weight", _weight())
+    close_errors: list[BaseException] = []
+
+    def close_output() -> None:
+        try:
+            output.close()
+        except BaseException as exc:
+            close_errors.append(exc)
+        finally:
+            closed.set()
+
+    close_thread = threading.Thread(target=close_output)
+    close_thread.start()
+    try:
+        assert started.wait(timeout=2)
+        assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it"
+    finally:
+        release.set()
+        close_thread.join(timeout=5)
+
+    assert not close_errors
+    writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+
+def test_payload_fsync_failure_prevents_publication(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+
+    def failing_fsync(fd: int) -> None:
+        try:
+            raise OSError(errno.EIO, "injected fsync failure")
+        finally:
+            os.close(fd)
+
+    monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync)
+    with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+        output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+        output.write_tensor("layer.weight", _weight())
+
+    with pytest.raises(OSError, match="injected fsync failure"):
+        writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    assert not (writer.staging_dir / "manifest.json").exists()
+    assert not (writer.staging_dir / "READY.json").exists()
+
+
+def test_unordered_payload_fallback_hashes_files_in_parallel(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    writer = _direct_writer(tmp_path)
+    specs = (
+        TensorWriteSpec("a", (2,), torch.float32),
+        TensorWriteSpec("b", (2,), torch.float32),
+    )
+    for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True):
+        with writer.open_tensor_file(name, (spec,)) as output:
+            output.write_tensor(spec.name, torch.tensor([1.0, 2.0]))
+
+    original_scan = filesystem_writer_module._scan_payload
+    barrier = threading.Barrier(2)
+    thread_ids: set[int] = set()
+    thread_ids_lock = threading.Lock()
+
+    def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object:
+        with thread_ids_lock:
+            thread_ids.add(threading.get_ident())
+        barrier.wait(timeout=5)
+        return original_scan(path, file_specs)
+
+    monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan)
+    manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+    assert len(manifest.files) == 2
+    assert len(thread_ids) == 2
+
+
 class FakeProducer:
     def __init__(
         self,
@@ -1026,7 +1172,9 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path:
 
     free_root = tmp_path / "free-limit"
     free_root.mkdir(mode=0o700)
-    unavailable_reserve = shutil.disk_usage(free_root).free + 1
+    # Leave enough margin for filesystem accounting to move between the
+    # preflight measurement and the capacity check.
+    unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30)
     free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
     free_result = free_limited.get_or_build(
         BuildRequest(identity),

vllm_omni/config/omni_config.py

@@ -699,6 +699,7 @@ class _DiffusionConfigProjection:
     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
+    dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -756,6 +757,7 @@ class _DiffusionConfigProjection:
             TransformerConfig,
             build_attention_config,
             parse_kv_cache_skip_selector,
+            validate_dlo_host_registration_options,
             validate_host_weight_runtime_options,
         )
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
@@ -838,6 +840,12 @@ class _DiffusionConfigProjection:
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
         )
+        self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=self.dlo_host_registration_limit_gib,
+            enable_dlo=self.enable_distributed_layerwise_offload,
+            use_allgather=self.dlo_use_allgather,
+            hwr_mode=self.host_weight_runtime_mode,
+        )
 
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(

vllm_omni/config/stage_config.py

@@ -391,6 +391,7 @@ class StageDeployConfig:
     dlo_resident_layers: int | None = None
     host_weight_runtime_mode: str | None = None
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_limit_gib: float | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/diffusion/data.py

@@ -113,6 +113,22 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
         raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
 
 
+def validate_dlo_host_registration_options(
+    *,
+    limit_gib: object,
+    enable_dlo: bool,
+    use_allgather: bool,
+    hwr_mode: object,
+) -> float:
+    """Validate the optional transport budget without probing CUDA or HWR."""
+    value = float(limit_gib)
+    if not math.isfinite(value) or value < 0:
+        raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
+    if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
+        raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime")
+    return value
+
+
 def parse_kv_cache_skip_selector(
     selector: str | list[int] | tuple[int, ...] | set[int] | None,
 ) -> set[int] | None:
@@ -784,6 +800,9 @@ class OmniDiffusionConfig:
     # existing loader/storage path.
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    # Optional per-worker ceiling for registering final-layout HWR mappings.
+    # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
+    dlo_host_registration_limit_gib: float = 0.0
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
@@ -1164,6 +1183,12 @@ class OmniDiffusionConfig:
             mode=self.host_weight_runtime_mode,
             root=self.host_weight_runtime_root,
         )
+        self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=self.dlo_host_registration_limit_gib,
+            enable_dlo=self.enable_distributed_layerwise_offload,
+            use_allgather=self.dlo_use_allgather,
+            hwr_mode=self.host_weight_runtime_mode,
+        )
 
         if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
             raise ValueError(

vllm_omni/diffusion/model_loader/host_weight_loader.py

@@ -27,7 +27,10 @@ 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
+    from vllm_omni.diffusion.model_loader.host_weights.source_identity import (
+        NodeSourceDigestCache,
+        PreparedWeightSource,
+    )
 
 
 class _HWRCommitError(RuntimeError):
@@ -215,6 +218,7 @@ class HWRLoaderMixin:
         *,
         load_format: str,
         sources: Sequence[object],
+        source_digest_cache: NodeSourceDigestCache | None = None,
     ) -> FinalLayoutIdentityContext:
         from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
 
@@ -293,6 +297,7 @@ class HWRLoaderMixin:
             prepared_sources=prepared_sources,
             request=request,
             policy=FINAL_LAYOUT_BF16_POLICY,
+            source_digest_cache=source_digest_cache,
         )
 
     def _resolve_hwr(
@@ -306,7 +311,10 @@ class HWRLoaderMixin:
         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.diffusion.model_loader.host_weights import (
+            FinalLayoutTensorRestorer,
+            NodeSourceDigestCache,
+        )
         from vllm_omni.host_weight_runtime import (
             HostWeightLeaseCarrier,
             HostWeightRuntime,
@@ -347,7 +355,6 @@ class HWRLoaderMixin:
                 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")
@@ -362,6 +369,17 @@ class HWRLoaderMixin:
                 ),
             )
         )
+        source_digest_cache = NodeSourceDigestCache(
+            root,
+            timeout_seconds=runtime.config.wait.coordination_timeout_seconds,
+        )
+        context = self._build_hwr_context(
+            model,
+            modules,
+            load_format=load_format,
+            sources=sources,
+            source_digest_cache=source_digest_cache,
+        )
         resolution = runtime.resolve(context.identity)
         state: dict[str, object] = {
             "mode": mode,

vllm_omni/diffusion/model_loader/host_weights/__init__.py

@@ -19,7 +19,7 @@ from .producers import (
     FinalLayoutBF16Producer,
 )
 from .restorer import FinalLayoutTensorRestorer
-from .source_identity import PreparedWeightSource, WeightSourceKind
+from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind
 
 __all__ = [
     "FINAL_LAYOUT_BF16_POLICY",
@@ -32,6 +32,7 @@ __all__ = [
     "FinalLayoutRequest",
     "FinalLayoutTensorRestorer",
     "ImplementationIdentity",
+    "NodeSourceDigestCache",
     "PreparedWeightSource",
     "WeightSourceKind",
     "build_final_layout_identity",

vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py

@@ -27,6 +27,7 @@ from .contracts import (
 )
 from .source_identity import (
     FinalLayoutSourceContext,
+    NodeSourceDigestCache,
     PreparedWeightSource,
     resolve_final_layout_source_identity,
 )
@@ -78,6 +79,7 @@ def build_final_layout_identity(
     prepared_sources: Sequence[PreparedWeightSource],
     request: FinalLayoutRequest,
     policy: FinalLayoutTensorPolicy,
+    source_digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutIdentityContext:
     """Build exact source, representation, layout, and implementation identity."""
     if not isinstance(request, FinalLayoutRequest):
@@ -99,6 +101,7 @@ def build_final_layout_identity(
         prepared_sources,
         model_id=request.model_id,
         target_names=target_names,
+        digest_cache=source_digest_cache,
     )
     dit_names = tuple(name for name, _ in dit_modules)
     model_contracts = tuple(binding.contract for binding in bindings)

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

@@ -286,7 +286,7 @@ class FinalLayoutBF16Producer:
                 for record in shard
             )
             file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
-            with writer.open_tensor_file(file_name, specs) as output:
+            with writer.open_tensor_file(file_name, specs, ordered=True) as output:
                 for record in shard:
                     output.write_tensor(record.name, record.tensor.detach())
 

vllm_omni/diffusion/model_loader/host_weights/source_identity.py

@@ -5,9 +5,13 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import os
 import stat
+import time
+import uuid
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -15,6 +19,7 @@ from pathlib import Path
 import regex as re
 
 from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError
 from vllm_omni.host_weight_runtime.identity import canonical_json
 
 from .contracts import (
@@ -101,6 +106,140 @@ class _FileSnapshot:
         )
 
 
+@dataclass(frozen=True)
+class _FileState:
+    size: int
+    device: int
+    inode: int
+    mtime_ns: int
+    ctime_ns: int
+    symlink_target: str | None
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "size": self.size,
+            "device": self.device,
+            "inode": self.inode,
+            "mtime_ns": self.mtime_ns,
+            "ctime_ns": self.ctime_ns,
+            "symlink_target": self.symlink_target,
+        }
+
+
+def _observe(path: Path) -> _FileState:
+    try:
+        current = path.stat()
+    except OSError as exc:
+        raise ValueError(f"cannot stat canonical weight file {path}") from exc
+    if not stat.S_ISREG(current.st_mode):
+        raise ValueError(f"canonical weight source is not a regular file: {path}")
+    return _FileState(
+        size=current.st_size,
+        device=current.st_dev,
+        inode=current.st_ino,
+        mtime_ns=current.st_mtime_ns,
+        ctime_ns=current.st_ctime_ns,
+        symlink_target=os.readlink(path) if path.is_symlink() else None,
+    )
+
+
+def _sha256_file(path: Path, expected: _FileState) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        before = os.fstat(handle.fileno())
+        while chunk := handle.read(_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+        after = os.fstat(handle.fileno())
+    observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns)
+    completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns)
+    expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns)
+    if observed != expected_stat or completed != expected_stat or _observe(path) != expected:
+        raise FinalLayoutContractError(
+            FinalLayoutContractCode.SOURCE_CHANGED,
+            f"canonical weight file changed while hashing: {path}",
+        )
+    return digest.hexdigest()
+
+
+class NodeSourceDigestCache:
+    """Share content-backed local-source fingerprints across node workers."""
+
+    def __init__(self, root: Path, *, timeout_seconds: float) -> None:
+        self.entries = root / "source-digests-v1" / "entries"
+        self.locks = root / "source-digests-v1" / "locks"
+        self.timeout_seconds = timeout_seconds
+        self._available = True
+        try:
+            self.entries.mkdir(parents=True, exist_ok=True, mode=0o700)
+            self.locks.mkdir(parents=True, exist_ok=True, mode=0o700)
+        except OSError:
+            self._available = False
+
+    @staticmethod
+    def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None:
+        try:
+            cached = json.loads(entry.read_bytes())
+            if not isinstance(cached, dict):
+                return None
+            cached_state = cached["state"]
+            content_id = cached["content_id"]
+            record_sha256 = cached["record_sha256"]
+        except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError):
+            return None
+        record = {
+            "content_id": content_id,
+            "schema_version": cached.get("schema_version"),
+            "state": cached_state,
+        }
+        if (
+            record["schema_version"] != 1
+            or cached_state != state.to_dict()
+            or not isinstance(content_id, str)
+            or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None
+            or not isinstance(record_sha256, str)
+            or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256
+            or _observe(path) != state
+        ):
+            return None
+        return content_id
+
+    def content_id(self, path: Path, state: _FileState) -> str:
+        if not self._available:
+            return f"sha256:{_sha256_file(path, state)}"
+        key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest()
+        entry = self.entries / f"{key}.json"
+        content_id: str | None = None
+        try:
+            with FileLock(
+                self.locks / f"{key}.lock",
+                exclusive=True,
+                deadline=time.monotonic() + self.timeout_seconds,
+            ):
+                if (cached := self._cached_content_id(entry, path, state)) is not None:
+                    return cached
+                content_id = f"sha256:{_sha256_file(path, state)}"
+                record: dict[str, object] = {
+                    "content_id": content_id,
+                    "schema_version": 1,
+                    "state": state.to_dict(),
+                }
+                record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest()
+                temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+                try:
+                    temporary.write_bytes(canonical_json(record))
+                    os.chmod(temporary, 0o600)
+                    os.replace(temporary, entry)
+                finally:
+                    temporary.unlink(missing_ok=True)
+                return content_id
+        except (FileLockTimeoutError, OSError):
+            # Digest sharing is an optimization. Coordination or cache I/O
+            # failures must not make a readable canonical source unavailable.
+            if content_id is not None and _observe(path) == state:
+                return content_id
+            return f"sha256:{_sha256_file(path, state)}"
+
+
 @dataclass(frozen=True)
 class _SourceSnapshot:
     semantic: CanonicalJson
@@ -129,14 +268,6 @@ class FinalLayoutSourceContext:
             )
 
 
-def _sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        while chunk := handle.read(_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
 def _hf_blob_content_id(
     path: Path,
     *,
@@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str:
     return value
 
 
-def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
+def _snapshot_source(
+    source: PreparedWeightSource,
+    digest_cache: NodeSourceDigestCache | None,
+) -> _SourceSnapshot:
     root = source.resolved_root.expanduser().resolve()
     if not root.is_dir():
         raise ValueError(f"canonical source root is not a directory: {root}")
@@ -210,8 +344,7 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
         if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None:
             immutable_revision = requested.lower()
 
-    files: list[_FileSnapshot] = []
-    for path in sorted(source.weight_files, key=lambda item: str(item)):
+    def snapshot_file(path: Path) -> _FileSnapshot:
         candidate = path.expanduser()
         if not candidate.is_absolute():
             candidate = root / candidate
@@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot:
             relative_name = candidate.relative_to(root).as_posix()
         except ValueError as exc:
             raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc
-        try:
-            current = candidate.stat()
-        except OSError as exc:
-            raise ValueError(f"cannot stat canonical weight file {candidate}") from exc
-        if not stat.S_ISREG(current.st_mode):
-            raise ValueError(f"canonical weight source is not a regular file: {candidate}")
-        symlink_target = os.readlink(candidate) if candidate.is_symlink() else None
+        current = _observe(candidate)
         content_id = _hf_blob_content_id(
             candidate,
             source=source,
             source_root=root,
         )
         if content_id is None:
-            content_id = f"sha256:{_sha256_file(candidate)}"
-        files.append(
-            _FileSnapshot(
-                path=candidate,
-                relative_name=relative_name,
-                size=current.st_size,
-                device=current.st_dev,
-                inode=current.st_ino,
-                mtime_ns=current.st_mtime_ns,
-                ctime_ns=current.st_ctime_ns,
-                symlink_target=symlink_target,
-                content_id=content_id,
+            content_id = (
+                digest_cache.content_id(candidate, current)
+                if digest_cache is not None
+                else f"sha256:{_sha256_file(candidate, current)}"
             )
+        return _FileSnapshot(
+            path=candidate,
+            relative_name=relative_name,
+            size=current.size,
+            device=current.device,
+            inode=current.inode,
+            mtime_ns=current.mtime_ns,
+            ctime_ns=current.ctime_ns,
+            symlink_target=current.symlink_target,
+            content_id=content_id,
         )
 
+    paths = sorted(source.weight_files, key=lambda item: str(item))
+    with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor:
+        files = list(executor.map(snapshot_file, paths))
+
     file_semantics = [file.semantic_dict() for file in files]
     file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest()
     revision = immutable_revision or f"content-{file_fingerprint}"
@@ -306,11 +439,12 @@ def resolve_final_layout_source_identity(
     *,
     model_id: str,
     target_names: frozenset[str],
+    digest_cache: NodeSourceDigestCache | None = None,
 ) -> FinalLayoutSourceContext:
     """Resolve exact source identity for files covering the owned tensors."""
     bindings = _resolve_target_sources(prepared_sources, target_names)
     selected_sources = frozenset(bindings.values())
-    snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources}
+    snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources}
     snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded))
     source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots]
     target_bindings = [
@@ -345,6 +479,7 @@ def resolve_final_layout_source_identity(
 
 __all__ = [
     "FinalLayoutSourceContext",
+    "NodeSourceDigestCache",
     "PreparedWeightSource",
     "WeightSourceKind",
     "resolve_final_layout_source_identity",

vllm_omni/diffusion/offloader/base.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
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -10,7 +10,7 @@ import torch
 from torch import nn
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options
 
 logger = init_logger(__name__)
 
@@ -52,6 +52,9 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Optional per-worker ceiling for registering an HWR mmap. Zero means no
+    # additional ceiling; pin_cpu_memory controls whether registration is tried.
+    dlo_host_registration_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -122,6 +125,12 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+            limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
+            enable_dlo=enable_distributed_layerwise_offload,
+            use_allgather=dlo_use_allgather,
+            hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"),
+        )
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -157,6 +166,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA implementation of read-only host-mapping registration."""
+
+from __future__ import annotations
+
+import ctypes
+import mmap
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+from .host_registration import (
+    HostRegistrationBudgetError,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+)
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
+
+
+class _CudaRuntime(Protocol):
+    def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ...
+
+    def cudaHostUnregister(self, address: int) -> int: ...
+
+    def cudaGetErrorString(self, error: int) -> str | bytes: ...
+
+    def cudaGetLastError(self) -> int: ...
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _mapped_regions(regions: Sequence[MappedHostRegion]) -> tuple[_AddressRange, ...]:
+    """Resolve page spans without merging unrelated artifact files."""
+    by_file: dict[str, list[tuple[int, int]]] = defaultdict(list)
+    for region in regions:
+        if not region.read_only:
+            raise HostRegistrationError(
+                f"CUDA direct H2D requires an immutable mapping, but {region.file_name!r} is writable"
+            )
+        by_file[region.file_name].append((region.address, region.length))
+    return tuple(span for ranges in by_file.values() for span in _coalesce_ranges(ranges))
+
+
+def _error_message(runtime: _CudaRuntime, error: int) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+def _consume_last_cuda_error(runtime: _CudaRuntime, expected_error: int) -> None:
+    """Clear one handled CUDA Runtime error before returning to PyTorch."""
+    try:
+        get_last_error = getattr(runtime, "cudaGetLastError", None)
+        if get_last_error is None:
+            get_last_error = ctypes.CDLL(None).cudaGetLastError
+            get_last_error.argtypes = []
+            get_last_error.restype = ctypes.c_int
+        pending_error = int(get_last_error())
+    except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc:
+        raise HostRegistrationError("cannot clear CUDA's pending error after host-registration failure") from exc
+
+    if pending_error not in (0, expected_error):
+        raise HostRegistrationError(
+            f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
+        )
+
+
+def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
+    error_code = int(error)
+    message = _error_message(runtime, error_code)
+    try:
+        _consume_last_cuda_error(runtime, error_code)
+    except HostRegistrationError as exc:
+        raise HostRegistrationError(f"{message}; {exc}") from exc
+    return message
+
+
+def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
+    """Query support required by immutable artifact mappings."""
+    try:
+        get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None)
+        if get_attribute is None:
+            get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute
+            get_attribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
+            get_attribute.restype = ctypes.c_int
+    except (AttributeError, OSError) as exc:
+        raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc
+
+    supported = ctypes.c_int()
+    try:
+        error = get_attribute(
+            ctypes.byref(supported),
+            _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED,
+            torch.accelerator.current_device_index(),
+        )
+    except Exception as exc:
+        raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc
+    if int(error) != 0:
+        raise HostRegistrationError(
+            f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}"
+        )
+    return bool(supported.value)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for immutable file-backed CPU regions."""
+
+    def __init__(self, runtime: _CudaRuntime, regions: tuple[_AddressRange, ...]) -> None:
+        self._runtime = runtime
+        self._regions = regions
+
+    @classmethod
+    def create(
+        cls,
+        regions: Sequence[MappedHostRegion],
+        *,
+        max_bytes: int | None,
+    ) -> CudaHostRegistration:
+        mapped = _mapped_regions(regions)
+        total_bytes = sum(region.size for region in mapped)
+        if max_bytes is not None and total_bytes > max_bytes:
+            raise HostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not mapped:
+            raise HostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise HostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+        if not _supports_read_only_host_registration(runtime):
+            raise HostRegistrationError(
+                "CUDA device does not support read-only host registration required by immutable HWR mappings"
+            )
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in mapped:
+                error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY)
+                if int(error) != 0:
+                    raise HostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            rollback_failed: list[_AddressRange] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
+                        )
+                        rollback_failed.append(region)
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+                    rollback_failed.append(region)
+            if rollback_errors:
+                active_registration = cls(runtime, tuple(reversed(rollback_failed)))
+                raise HostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}",
+                    active_registration=active_registration,
+                ) from exc
+            if isinstance(exc, HostRegistrationError):
+                raise
+            raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, mapped)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> tuple[str, ...]:
+        """Unregister every range and retain failures for a later retry."""
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        return tuple(errors)

vllm_omni/diffusion/offloader/distributed_layerwise_backend.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
 """Distributed Layerwise Offload backend with double-buffered H2D.
 
 This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
@@ -18,6 +18,8 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import threading
+import time
 from itertools import chain
 from typing import Any
 
@@ -35,6 +37,12 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
+from .host_registration import (
+    HostRegistration,
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+    register_host_mappings,
+)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -51,6 +59,38 @@ from .tensor_utils import (
 
 logger = init_logger(__name__)
 
+# A backend normally owns both objects below. This process-lifetime safety
+# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup
+# failure unwinds startup and the backend itself becomes unreachable. A clean
+# retry removes the pair before closing the lease.
+_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = []
+_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock()
+
+
+def _retain_active_hwr_registration(
+    registration: HostRegistration,
+    lease: HostWeightLease,
+) -> None:
+    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
+        if not any(
+            candidate is registration and candidate_lease is lease
+            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
+        ):
+            _ACTIVE_HWR_REGISTRATIONS.append((registration, lease))
+
+
+def _forget_active_hwr_registration(
+    registration: HostRegistration,
+    lease: HostWeightLease,
+) -> None:
+    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
+        _ACTIVE_HWR_REGISTRATIONS[:] = [
+            (candidate, candidate_lease)
+            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
+            if candidate is not registration or candidate_lease is not lease
+        ]
+
+
 # Threshold (in MB) for deciding whether a non-block DiT submodule should
 # use layerwise offload (streaming hooks) or be moved to GPU as a resident
 # module.  Submodules larger than this are offloaded to save HBM; smaller
@@ -94,6 +134,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -445,6 +486,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -457,18 +518,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -500,6 +550,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
         return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     @torch.compiler.disable
     def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
         """Prepare next block's weights into the shared device buffer for *slot*.
@@ -515,17 +591,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-            with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
-                    gw = gpu_weights[dtype]
-                    async_copy = non_blocking and cpu_shard.is_pinned()
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                evt.record(self.copy_stream)
-            if self.rank_local_mmap:
-                # The CPU slot may be overwritten only after this H2D copy has
-                # finished.  The shared event protects reuse by another hook.
-                self.cpu_staging_events[slot] = evt
+            if self.rank_local_mmap and self.registered_mmap:
+                with current_omni_platform.stream(self.copy_stream):
+                    self._copy_mmap_sources_to_device(
+                        self.cpu_sources,
+                        self.metadata,
+                        gpu_weights,
+                        non_blocking=non_blocking,
+                    )
+                    evt.record(self.copy_stream)
+            else:
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -735,6 +821,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -805,6 +892,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -902,8 +997,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._host_weight_lease: HostWeightLease | None = None
+        self._host_registration: HostRegistration | None = None
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
 
     def load_resident_layers(self) -> None:
@@ -916,6 +1013,136 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if self._resident_layer_group is not None:
             self._resident_layer_group.offload()
 
+    def _rank_local_source_tensors(
+        self,
+        hooks: list[DistributedLayerwiseOffloadHook],
+    ) -> tuple[torch.Tensor, ...]:
+        """Return the exact CPU sources that direct H2D would consume."""
+        tensors: list[torch.Tensor] = []
+        seen: set[int] = set()
+
+        def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None:
+            for sources in cpu_sources.values():
+                for source in sources:
+                    tensor = source["tensor"]
+                    if id(tensor) not in seen:
+                        seen.add(id(tensor))
+                        tensors.append(tensor)
+
+        for hook in hooks:
+            collect(hook.cpu_sources)
+        if self._resident_layer_group is not None:
+            for state in self._resident_layer_group._states:
+                collect(state["cpu_sources"])
+        return tuple(tensors)
+
+    def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool:
+        """Register the complete final-layout lease under pinned-memory policy."""
+        lease = self._host_weight_lease
+        if lease is None:
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
+            return False
+        if not lease.mapped_regions or not source_tensors:
+            logger.warning("HWR mmap registration found no mapped sources; using bounded host staging")
+            return False
+
+        limit_gib = self.config.dlo_host_registration_limit_gib
+        max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
+        started = time.perf_counter()
+        try:
+            registration = register_host_mappings(
+                lease.mapped_regions,
+                device=self.device,
+                max_bytes=max_bytes,
+            )
+            try:
+                unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()]
+            except Exception as exc:
+                errors = registration.close()
+                if errors:
+                    self._host_registration = registration
+                    _retain_active_hwr_registration(registration, lease)
+                    raise HostRegistrationCleanupError(
+                        "CUDA registration succeeded but pinned-source verification failed, "
+                        f"and rollback failed: {errors[:3]}"
+                    ) from exc
+                raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc
+            if unpinned:
+                errors = registration.close()
+                if errors:
+                    self._host_registration = registration
+                    _retain_active_hwr_registration(registration, lease)
+                    raise HostRegistrationCleanupError(
+                        "CUDA registration succeeded but PyTorch rejected mapped sources, "
+                        f"and rollback failed: {errors[:3]}"
+                    )
+                raise HostRegistrationError(
+                    "CUDA registration succeeded but PyTorch did not recognize "
+                    f"{len(unpinned)} mapped source(s) as pinned"
+                )
+        except HostRegistrationCleanupError as exc:
+            # Falling back could close a lease while the platform still owns
+            # one of its mappings. Fail startup and retain ownership for retry.
+            active_registration = exc.active_registration
+            if active_registration is not None:
+                self._host_registration = active_registration
+                _retain_active_hwr_registration(active_registration, lease)
+            logger.exception("HWR mmap registration rollback failed")
+            raise
+        except HostRegistrationError as exc:
+            logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc)
+            return False
+
+        self._host_registration = registration
+        _retain_active_hwr_registration(registration, lease)
+        logger.info(
+            "Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None:
+        """Select registered direct H2D or bounded staging once per backend."""
+        plan = self.host_weight_plan
+        if (
+            (not hooks and self._resident_layer_group is None)
+            or not self._using_rank_local_mmap
+            or plan is None
+            or plan.backing_kind != "host_weight_runtime"
+        ):
+            return
+
+        source_tensors = self._rank_local_source_tensors(hooks)
+        self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors)
+        for hook in hooks:
+            hook.registered_mmap = self._using_registered_mmap
+        if self._resident_layer_group is not None:
+            self._resident_layer_group.registered_mmap = self._using_registered_mmap
+            if self._using_registered_mmap:
+                self._resident_layer_group._cpu_staging_buffers.clear()
+
+    def _release_registered_mmap(self) -> None:
+        """Release every platform registration before closing the HWR lease."""
+        registration = self._host_registration
+        if registration is None:
+            return
+        errors = registration.close()
+        if errors:
+            lease = self._host_weight_lease
+            if lease is not None and not lease.closed:
+                _retain_active_hwr_registration(registration, lease)
+            logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3])
+            raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)")
+        lease = self._host_weight_lease
+        if lease is not None:
+            _forget_active_hwr_registration(registration, lease)
+        self._host_registration = None
+        logger.info("Unregistered HWR mmap ranges")
+
     def _load_weights_via_mmap(
         self,
         pipeline: nn.Module,
@@ -1465,11 +1692,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     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.
+                # sources; transport setup below selects registered direct
+                # H2D or the bounded two-slot staging fallback.
                 self._using_rank_local_mmap = True
                 logger.info(
-                    "DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging",
+                    "DLO consuming final-layout Host Weight Runtime lease %s",
                     self._host_weight_lease.provenance.resolution_id,
                 )
             elif self.host_weight_plan.backing_kind == "checkpoint_mmap":
@@ -1483,7 +1710,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             if self._using_rank_local_mmap:
                 logger.info(
                     "DLO rank-local host storage enabled: source pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "node-shared; transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1653,6 +1880,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             )
             pipeline._dlo_residency_controller = self
 
+        all_hooks = [hook for group in self._all_hook_groups for hook in group]
+        self._configure_hwr_transfer(all_hooks)
+
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
             if self._using_mmap and not self.enabled:
@@ -1662,17 +1892,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers
         # sized to the max block across ALL module groups (gen_layers +
         # language_model).  Groups execute sequentially, so 2 buffers suffice.
-        all_hooks: list[DistributedLayerwiseOffloadHook] = []
-        for group in self._all_hook_groups:
-            all_hooks.extend(group)
-
         unified_buffers = self._allocate_shared_buffers(all_hooks)
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1753,6 +1979,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def _release_mmap_handles(self) -> None:
         """Release source handles and the transport-owned HWR lease."""
+        if self._host_registration is not None:
+            raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active")
         self._mmap_transforms_by_tensor_id.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
@@ -1781,6 +2009,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def disable(self) -> None:
         has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
+        has_registration = self._host_registration is not None
         has_carrier = (
             self.host_weight_plan is not None
             and self.host_weight_plan.lease_carrier is not None
@@ -1796,12 +2025,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             not self.enabled
             and not hasattr(self, "_mmap_file_cache")
             and not has_open_lease
+            and not has_registration
             and not has_carrier
             and not has_partial_hooks
         ):
             return
 
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap or has_registration:
             current_omni_platform.synchronize()
 
         for blocks in self._blocks:
@@ -1818,9 +2048,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
+        self._release_registered_mmap()
         self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/offloader/host_registration.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Platform-neutral lifecycle for registering existing host mappings."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol
+
+import torch
+
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+
+class HostRegistrationError(RuntimeError):
+    """Raised when existing host ranges cannot be registered safely."""
+
+
+class HostRegistrationBudgetError(HostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class HostRegistrationCleanupError(HostRegistrationError):
+    """Raised when registered ranges cannot be released safely."""
+
+    def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None:
+        super().__init__(message)
+        self.active_registration = active_registration
+
+
+class HostRegistration(Protocol):
+    """Transport-owned registration for one process-local host mapping."""
+
+    @property
+    def total_bytes(self) -> int: ...
+
+    @property
+    def region_count(self) -> int: ...
+
+    def close(self) -> tuple[str, ...]: ...
+
+
+def register_host_mappings(
+    regions: Sequence[MappedHostRegion],
+    *,
+    device: torch.device,
+    max_bytes: int | None,
+) -> HostRegistration:
+    """Register immutable ranges with the active platform implementation."""
+    if device.type == "cuda":
+        from .cuda_host_registration import CudaHostRegistration
+
+        return CudaHostRegistration.create(regions, max_bytes=max_bytes)
+    raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}")
+
+
+__all__ = [
+    "HostRegistration",
+    "HostRegistrationBudgetError",
+    "HostRegistrationCleanupError",
+    "HostRegistrationError",
+    "register_host_mappings",
+]

vllm_omni/diffusion/worker/diffusion_worker.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
 
 """
 Diffusion Worker for vLLM-Omni.
@@ -897,11 +897,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-            if mgr is not None:
-                mgr.shutdown_prefetch()
-        destroy_distributed_env()
+        try:
+            if self.model_runner is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                try:
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -571,6 +571,7 @@ class OrchestratorArgs:
     dlo_resident_layers: int = 0
     host_weight_runtime_mode: str = "disabled"
     host_weight_runtime_root: str | None = None
+    dlo_host_registration_limit_gib: float = 0.0
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1061,6 +1061,7 @@ class AsyncOmniEngine:
             "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"),
+            "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
             "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

@@ -760,6 +760,16 @@ class OmniServeCommand(CLISubcommand):
                 "the same persistent path for population and serving."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-host-registration-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. "
+                "Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration "
+                "under the existing pinned-memory policy and fall back to bounded staging when unavailable."
+            ),
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",

vllm_omni/host_weight_runtime/filesystem/writer.py

@@ -12,17 +12,19 @@ import os
 import struct
 import uuid
 from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
 from types import TracebackType
 
 import torch
-from safetensors import safe_open
 
 from ..identity import WeightArtifactIdentity, canonical_json
 from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest
 from ..protocols import ProductionMetadata, TensorWriteSpec
 
 _HASH_CHUNK_BYTES = 8 * 1024**2
+_MAX_HASH_WORKERS = 8
 _MANIFEST_FILE = "manifest.json"
 _READY_FILE = "READY.json"
 
@@ -79,21 +81,47 @@ def _fsync_directory(path: Path) -> None:
         os.close(fd)
 
 
-def _file_sha256(path: Path) -> str:
-    digest = hashlib.sha256()
+def _fsync_payload(fd: int) -> None:
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+@dataclass(frozen=True)
+class _PayloadDigests:
+    file: str
+    tensors: dict[str, str]
+
+
+def _scan_payload(path: Path, specs: tuple[TensorWriteSpec, ...]) -> _PayloadDigests:
+    """Calculate file and tensor digests in one ordered read."""
+    file_digest = hashlib.sha256()
+    tensor_digests: dict[str, str] = {}
     with path.open("rb") as handle:
-        while chunk := handle.read(_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _tensor_sha256(tensor: torch.Tensor) -> str:
-    digest = hashlib.sha256()
-    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
-    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
-        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
-        digest.update(memoryview(chunk.numpy()))
-    return digest.hexdigest()
+        prefix = handle.read(8)
+        if len(prefix) != 8:
+            raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
+        header_size = struct.unpack("<Q", prefix)[0]
+        header = handle.read(header_size)
+        if len(header) != header_size:
+            raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}")
+        file_digest.update(prefix)
+        file_digest.update(header)
+        for spec in specs:
+            digest = hashlib.sha256()
+            remaining = spec.nbytes
+            while remaining:
+                chunk = handle.read(min(remaining, _HASH_CHUNK_BYTES))
+                if not chunk:
+                    raise ArtifactWriterError(f"truncated tensor {spec.name!r} in {path.name!r}")
+                file_digest.update(chunk)
+                digest.update(chunk)
+                remaining -= len(chunk)
+            tensor_digests[spec.name] = digest.hexdigest()
+        if handle.read(1):
+            raise ArtifactWriterError(f"unexpected trailing bytes in {path.name!r}")
+    return _PayloadDigests(file_digest.hexdigest(), tensor_digests)
 
 
 class _Coverage:
@@ -130,6 +158,8 @@ class StreamingSafeTensorFile:
         owner: FilesystemArtifactWriter,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool,
     ) -> None:
         self._owner = owner
         self.relative_name = _safe_relative_name(relative_name)
@@ -144,6 +174,7 @@ class StreamingSafeTensorFile:
         self.specs = tuple(sorted(specs, key=lambda spec: spec.name))
         self._specs = {spec.name: spec for spec in self.specs}
         self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs}
+        self._ordered = ordered
         self._closed = False
 
         offsets: dict[str, int] = {}
@@ -167,6 +198,7 @@ class StreamingSafeTensorFile:
         total_size = self._payload_offset + cursor
         owner._reserve_file(total_size)
         self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600)
+        prefix = struct.pack("<Q", len(encoded_header))
         try:
             if hasattr(os, "posix_fallocate"):
                 try:
@@ -177,13 +209,16 @@ class StreamingSafeTensorFile:
                     os.ftruncate(self._fd, total_size)
             else:  # pragma: no cover - Linux production path has posix_fallocate
                 os.ftruncate(self._fd, total_size)
-            _pwrite_all(self._fd, struct.pack("<Q", len(encoded_header)), 0)
+            _pwrite_all(self._fd, prefix, 0)
             _pwrite_all(self._fd, encoded_header, 8)
         except BaseException:
             os.close(self._fd)
             self._path.unlink(missing_ok=True)
             owner._forget_reserved_file(total_size)
             raise
+        self._file_digest = hashlib.sha256(prefix + encoded_header) if ordered else None
+        self._tensor_digests = {spec.name: hashlib.sha256() for spec in self.specs} if ordered else {}
+        self._write_cursor = self._payload_offset
         owner._opened(self)
 
     @staticmethod
@@ -226,9 +261,19 @@ class StreamingSafeTensorFile:
             raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc
 
     def _write(self, name: str, byte_offset: int, data: memoryview) -> None:
+        file_offset = self._payload_offset + self._offsets[name] + byte_offset
+        if self._ordered and file_offset != self._write_cursor:
+            raise ArtifactWriterError(
+                f"ordered tensor payload expected offset {self._write_cursor}, got {file_offset} for {name!r}"
+            )
         coverage = self._coverage[name]
         coverage.add(byte_offset, byte_offset + len(data))
-        _pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset)
+        _pwrite_all(self._fd, data, file_offset)
+        if self._ordered:
+            assert self._file_digest is not None
+            self._file_digest.update(data)
+            self._tensor_digests[name].update(data)
+            self._write_cursor += len(data)
 
     def close(self) -> None:
         if self._closed:
@@ -238,11 +283,19 @@ class StreamingSafeTensorFile:
             self.abort()
             raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}")
         self._closed = True
+        digests = None
+        if self._ordered:
+            assert self._file_digest is not None
+            digests = _PayloadDigests(
+                self._file_digest.hexdigest(),
+                {name: digest.hexdigest() for name, digest in self._tensor_digests.items()},
+            )
         try:
-            os.fsync(self._fd)
-        finally:
+            os.fchmod(self._fd, 0o444)
+            self._owner._closed(self, digests, self._fd)
+        except BaseException:
             os.close(self._fd)
-        self._owner._closed(self)
+            raise
 
     def abort(self) -> None:
         if self._closed:
@@ -281,20 +334,25 @@ class FilesystemArtifactWriter:
         self._capacity_check = capacity_check
         self._active: set[StreamingSafeTensorFile] = set()
         self._files: dict[str, tuple[TensorWriteSpec, ...]] = {}
+        self._digests: dict[str, _PayloadDigests] = {}
         self._reserved_bytes = 0
         self._finalized = False
+        self._fsync_executor: ThreadPoolExecutor | None = None
+        self._fsync_futures: list[Future[None]] = []
 
     def open_tensor_file(
         self,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool = False,
     ) -> StreamingSafeTensorFile:
         if self._finalized:
             raise ArtifactWriterError("cannot add payloads after artifact finalization")
         relative_name = _safe_relative_name(relative_name)
         if relative_name in self._files or any(item.relative_name == relative_name for item in self._active):
             raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once")
-        return StreamingSafeTensorFile(self, relative_name, specs)
+        return StreamingSafeTensorFile(self, relative_name, specs, ordered=ordered)
 
     def _reserve_file(self, size: int) -> None:
         self._capacity_check(self._reserved_bytes + size, size)
@@ -306,9 +364,14 @@ class FilesystemArtifactWriter:
     def _opened(self, writer: StreamingSafeTensorFile) -> None:
         self._active.add(writer)
 
-    def _closed(self, writer: StreamingSafeTensorFile) -> None:
+    def _closed(self, writer: StreamingSafeTensorFile, digests: _PayloadDigests | None, fd: int) -> None:
         self._active.remove(writer)
         self._files[writer.relative_name] = writer.specs
+        if digests is not None:
+            self._digests[writer.relative_name] = digests
+        if self._fsync_executor is None:
+            self._fsync_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hwr-fsync")
+        self._fsync_futures.append(self._fsync_executor.submit(_fsync_payload, fd))
 
     def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None:
         self._active.discard(writer)
@@ -318,6 +381,19 @@ class FilesystemArtifactWriter:
         for writer in tuple(self._active):
             with contextlib.suppress(Exception):
                 writer.abort()
+        with contextlib.suppress(Exception):
+            self._wait_for_payload_fsync()
+
+    def _wait_for_payload_fsync(self) -> None:
+        if self._fsync_executor is None:
+            return
+        executor = self._fsync_executor
+        futures = tuple(self._fsync_futures)
+        self._fsync_executor = None
+        self._fsync_futures.clear()
+        executor.shutdown(wait=True)
+        for future in futures:
+            future.result()
 
     def finalize(
         self,
@@ -335,46 +411,50 @@ class FilesystemArtifactWriter:
         if metadata.restorer_schema != identity.producer.restorer_schema:
             raise ArtifactWriterError("production metadata does not match the identity's restorer schema")
 
+        self._wait_for_payload_fsync()
+
         tensor_entries: list[TensorManifestEntry] = []
         file_entries: list[FileManifestEntry] = []
         seen_tensors: set[str] = set()
-        for file_name in sorted(self._files):
+        file_names = sorted(self._files)
+        unhashed = [file_name for file_name in file_names if file_name not in self._digests]
+        if unhashed:
+            with ThreadPoolExecutor(max_workers=min(_MAX_HASH_WORKERS, len(unhashed))) as executor:
+                scans = executor.map(
+                    lambda file_name: _scan_payload(self.staging_dir / file_name, self._files[file_name]),
+                    unhashed,
+                )
+                self._digests.update(zip(unhashed, scans, strict=True))
+
+        for file_name in file_names:
             path = self.staging_dir / file_name
-            os.chmod(path, 0o444)
-            with path.open("rb") as handle:
-                os.fsync(handle.fileno())
+            specs = self._files[file_name]
+            digests = self._digests[file_name]
             file_entries.append(
                 FileManifestEntry(
                     file_name=file_name,
                     size=path.stat().st_size,
-                    sha256=_file_sha256(path),
+                    sha256=digests.file,
                 )
             )
-            expected = {spec.name: spec for spec in self._files[file_name]}
-            with safe_open(str(path), framework="pt", device="cpu") as handle:
-                if set(handle.keys()) != set(expected):
-                    raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}")
-                for name in sorted(expected):
-                    if name in seen_tensors:
-                        raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files")
-                    seen_tensors.add(name)
-                    spec = expected[name]
-                    tensor = handle.get_tensor(name)
-                    if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype:
-                        raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata")
-                    entry = TensorManifestEntry(
-                        name=name,
+            for spec in specs:
+                if spec.name in seen_tensors:
+                    raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files")
+                seen_tensors.add(spec.name)
+                tensor_entries.append(
+                    TensorManifestEntry(
+                        name=spec.name,
                         kind=spec.kind,
                         role=spec.role,
                         file_name=file_name,
-                        storage_key=name,
-                        shape=tuple(tensor.shape),
-                        dtype=str(tensor.dtype),
-                        stride=tuple(tensor.stride()),
-                        nbytes=tensor.numel() * tensor.element_size(),
-                        sha256=_tensor_sha256(tensor),
+                        storage_key=spec.name,
+                        shape=spec.shape,
+                        dtype=str(spec.dtype),
+                        stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()),
+                        nbytes=spec.nbytes,
+                        sha256=digests.tensors[spec.name],
                     )
-                    tensor_entries.append(entry)
+                )
 
         ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name))
         content_digest = hashlib.sha256()

vllm_omni/host_weight_runtime/protocols.py

@@ -88,6 +88,8 @@ class ArtifactWriter(Protocol):
         self,
         relative_name: str,
         specs: tuple[TensorWriteSpec, ...],
+        *,
+        ordered: bool = False,
     ) -> TensorFileWriter: ...