[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>

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