[Diffusion] Address runtime cache review feedback

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

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 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 (
@@ -15,6 +16,10 @@ from vllm_omni.diffusion.offloader.cuda_host_registration import (
     CudaHostRegistrationError,
     _coalesce_ranges,
 )
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationError,
+    register_host_mappings,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -78,6 +83,25 @@ class _FakeRuntime:
         return f"error-{error}"
 
 
+def test_platform_factory_reports_unsupported_registration() -> None:
+    with pytest.raises(HostRegistrationError, match="not supported on cpu"):
+        register_host_mappings({}, device=torch.device("cpu"), max_bytes=4096)
+
+
+def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    sources = {"weights": []}
+
+    def create(actual_sources, *, max_bytes):
+        assert actual_sources is sources
+        assert max_bytes == 4096
+        return sentinel
+
+    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+
+    assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel
+
+
 def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
     assert _coalesce_ranges(
         [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -1056,11 +1056,12 @@ class TestMmapWeightLoading:
 
         create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
 
-        def create(sources, *, max_bytes):
+        def register(sources, *, device, max_bytes):
+            assert device == torch.device("cuda")
             create_calls.append((sources, max_bytes))
             return Registration()
 
-        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -1084,6 +1085,24 @@ class TestMmapWeightLoading:
 
         assert events == ["unregister", "mmap"]
 
+    def test_runtime_cache_registration_falls_back_on_unsupported_platform(
+        self,
+        patched_offload_runtime,
+    ):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.0,
+            ),
+            torch.device("cpu"),
+        )
+        backend._runtime_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
+
+        assert not backend._try_register_runtime_cache_mmap()
+        assert backend._host_registration is None
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):