[Diffusion] Simplify DLO host plan setup

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

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -310,6 +310,7 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
 
     assert loader.load_model(load_device="cpu") is model
     assert not loaded_ordinary_weights
+    assert plan.requires_backend
     assert loader.take_host_weight_plan() is plan
     assert loader.take_host_weight_plan() is None
 
@@ -449,7 +450,6 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
         backing_kind="host_weight_cache",
         bindings={},
         runtime_layout_key="layout-key",
-        post_load_complete=True,
     )
 
     def load_weights(_model):
@@ -530,7 +530,6 @@ def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeyp
         backing_kind="host_weight_cache",
         bindings={},
         runtime_layout_key="layout-key",
-        post_load_complete=True,
     )
 
     loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]

tests/diffusion/model_loader/test_host_weight_cache.py

@@ -90,7 +90,7 @@ def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
     assert first.plan is not None
     assert second.plan is not None
     assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
-    assert first.plan.post_load_complete
+    assert not first.plan.requires_backend
     assert {binding.file_path for binding in first.plan.bindings.values()} == {
         binding.file_path for binding in second.plan.bindings.values()
     }

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -842,7 +842,6 @@ class TestMmapWeightLoading:
                 )
             },
             runtime_layout_key="runtime-layout",
-            post_load_complete=True,
             expected_file_digests={str(weight_file): file_digest(weight_file)},
         )
         backend = DistributedLayerwiseOffloadBackend(
@@ -881,7 +880,6 @@ class TestMmapWeightLoading:
                 )
             },
             runtime_layout_key="runtime-layout",
-            post_load_complete=True,
             expected_file_digests={str(tmp_path / "missing.safetensors"): "unused"},
         )
         backend = DistributedLayerwiseOffloadBackend(
@@ -926,7 +924,6 @@ class TestMmapWeightLoading:
                 )
             },
             runtime_layout_key="runtime-layout",
-            post_load_complete=True,
             expected_file_digests={str(weight_file): expected_digest},
         )
         backend = DistributedLayerwiseOffloadBackend(
@@ -1065,7 +1062,6 @@ class TestMmapWeightLoading:
             backing_kind="host_weight_cache",
             bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
             runtime_layout_key="runtime-layout",
-            post_load_complete=True,
             expected_file_digests={str(weight_file): file_digest(weight_file)},
         )
         backend = DistributedLayerwiseOffloadBackend(
@@ -1113,7 +1109,6 @@ class TestMmapWeightLoading:
             backing_kind="host_weight_cache",
             bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
             runtime_layout_key="runtime-layout",
-            post_load_complete=True,
             expected_file_digests={str(weight_file): "stale-digest"},
         )
         backend = DistributedLayerwiseOffloadBackend(
@@ -1269,6 +1264,47 @@ class TestMmapWeightLoading:
         assert not backend._try_register_host_weight_cache_mmap()
         assert backend._host_registration is None
 
+    def test_host_weight_cache_transfer_is_configured_once_for_all_consumers(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cuda"),
+            host_weight_plan=HostWeightPlan(
+                backing_kind="host_weight_cache",
+                bindings={},
+            ),
+        )
+        backend._using_rank_local_mmap = True
+        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
+        streamed_hook = SimpleNamespace(registered_mmap=False)
+        resident_group = SimpleNamespace(
+            registered_mmap=False,
+            _cpu_staging_buffers=[{"buffer": torch.ones(1)}],
+        )
+        backend._resident_layer_group = resident_group
+        register_calls: list[None] = []
+
+        def register() -> bool:
+            register_calls.append(None)
+            return True
+
+        monkeypatch.setattr(backend, "_try_register_host_weight_cache_mmap", register)
+
+        backend._configure_host_weight_cache_transfer([streamed_hook])  # type: ignore[list-item]
+
+        assert register_calls == [None]
+        assert streamed_hook.registered_mmap
+        assert resident_group.registered_mmap
+        assert resident_group._cpu_staging_buffers == []
+        assert backend._host_weight_cache_mapped_sources == {}
+
     def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
         self,
         patched_offload_runtime,

vllm_omni/diffusion/model_loader/host_weight_cache.py

@@ -401,7 +401,6 @@ def _validate_entry(
         backing_kind="host_weight_cache",
         bindings=bindings,
         runtime_layout_key=cache_key,
-        post_load_complete=True,
         expected_file_digests={
             str(entry_dir / filename): metadata["sha256"] for filename, metadata in file_manifest.items()
         },

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -9,6 +9,7 @@ import json
 import os
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
+from typing import Literal
 
 import torch
 from safetensors import safe_open
@@ -37,13 +38,17 @@ class TensorBinding:
 class HostWeightPlan:
     """Complete, prevalidated host backing consumed by an offload backend."""
 
-    backing_kind: str
+    backing_kind: Literal["checkpoint_mmap", "host_weight_cache"]
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
     runtime_layout_key: str | None = None
-    post_load_complete: bool = False
     expected_file_digests: dict[str, str] | None = None
 
+    @property
+    def requires_backend(self) -> bool:
+        """Whether ordinary loading was skipped and this plan must be consumed."""
+        return self.backing_kind == "checkpoint_mmap"
+
 
 @dataclass(frozen=True)
 class HostWeightPlanResult:

vllm_omni/diffusion/offloader/__init__.py

@@ -86,7 +86,7 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
-    plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete
+    plan_requires_backend = host_weight_plan is not None and host_weight_plan.requires_backend
     if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
         if plan_requires_backend:
             raise RuntimeError(

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1179,8 +1179,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         """Replace final ordinary-loader tensors with validated cache views."""
         from safetensors import safe_open
 
-        if not plan.post_load_complete:
-            raise RuntimeError("A host weight cache plan must describe fully post-processed weights")
         if any(binding.transform is not None for binding in plan.bindings.values()):
             raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms")
 
@@ -1307,6 +1305,29 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         )
         return True
 
+    def _configure_host_weight_cache_transfer(
+        self,
+        hooks: list[DistributedLayerwiseOffloadHook],
+    ) -> None:
+        """Select registered mmap or staging once for every cache consumer."""
+        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_cache"
+        ):
+            return
+
+        self._using_registered_mmap = self._try_register_host_weight_cache_mmap()
+        self._host_weight_cache_mapped_sources.clear()
+        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) -> bool:
         registration = self._host_registration
         if registration is None:
@@ -1866,19 +1887,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             )
             pipeline._dlo_residency_controller = self
 
+        all_hooks = [hook for group in self._all_hook_groups for hook in group]
+        self._configure_host_weight_cache_transfer(all_hooks)
+
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
-            if (
-                self.enabled
-                and self._using_rank_local_mmap
-                and self.host_weight_plan.backing_kind == "host_weight_cache"
-            ):
-                self._using_registered_mmap = self._try_register_host_weight_cache_mmap()
-                self._host_weight_cache_mapped_sources.clear()
-                assert 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()
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
             return
@@ -1886,23 +1899,11 @@ 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)
 
-        if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "host_weight_cache":
-            self._using_registered_mmap = self._try_register_host_weight_cache_mmap()
-            self._host_weight_cache_mapped_sources.clear()
-            for hook in all_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
-
         unified_cpu_staging = None
         cpu_staging_events = None
         if self._using_rank_local_mmap and not self._using_registered_mmap: