[Diffusion] Prefer host-cache registration when pinned

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

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -498,7 +498,8 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
     assert loader.take_host_weight_plan() is runtime_plan
 
 
-def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+@pytest.mark.parametrize("pin_cpu_memory", [True, False])
+def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
 
@@ -515,8 +516,10 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
         quantization_config=None,
         enable_distributed_layerwise_offload=True,
         dlo_use_allgather=False,
+        dlo_use_host_weight_cache=True,
+        pin_cpu_memory=pin_cpu_memory,
         dlo_host_weight_cache_dir=str(tmp_path),
-        dlo_host_weight_cache_pin_limit_gib=1.0,
+        dlo_host_weight_cache_pin_limit_gib=0.0,
         model="unused",
     )
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
@@ -536,7 +539,7 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
     loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
 
     def unexpected_checkpoint_plan(*_args, **_kwargs):
-        pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan")
+        pytest.fail("selected host weight cache mode must not select the raw checkpoint mmap plan")
 
     def build_cache(_pipeline, **_kwargs):
         events.append("cache")

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -119,6 +119,20 @@ def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> No
     assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel
 
 
+def test_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    sources = {"weights": []}
+
+    def create(actual_sources, *, max_bytes):
+        assert actual_sources is sources
+        assert max_bytes is None
+        return sentinel
+
+    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+
+    assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel
+
+
 def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
     assert _coalesce_ranges(
         [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
@@ -143,6 +157,21 @@ def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytes
     assert runtime.registered == []
 
 
+def test_registration_without_additional_budget_registers_complete_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+        max_bytes=None,
+    )
+
+    assert registration.total_bytes == 8192
+    assert len(runtime.registered) == 1
+    assert registration.close() == []
+
+
 def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0], read_only_registration_supported=False)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -1207,6 +1207,68 @@ class TestMmapWeightLoading:
 
         assert events == ["unregister", "mmap"]
 
+    def test_host_weight_cache_registration_defaults_to_pinned_memory_policy(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                return []
+
+        registration = Registration()
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = []
+
+        def register(sources, *, device, max_bytes):
+            assert device == torch.device("cuda")
+            create_calls.append((sources, max_bytes))
+            return registration
+
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_host_weight_cache_pin_limit_gib=0.0,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._host_weight_cache_mapped_sources = sources
+
+        assert backend._try_register_host_weight_cache_mmap()
+        assert create_calls == [(sources, None)]
+        assert backend._host_registration is registration
+
+    def test_host_weight_cache_registration_respects_disabled_pinned_memory(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        monkeypatch.setattr(
+            dist_backend_module,
+            "register_host_mappings",
+            lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"),
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                dlo_host_weight_cache_pin_limit_gib=0.0,
+            ),
+            torch.device("cuda"),
+        )
+        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
+
+        assert not backend._try_register_host_weight_cache_mmap()
+        assert backend._host_registration is None
+
     def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
         self,
         patched_offload_runtime,