[Diffusion] Gate CUDA read-only host registration

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

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -542,12 +542,15 @@ After mapping and validating every final tensor, DLO groups source storages by
 backing file, page-aligns their address ranges, and coalesces overlapping or
 adjacent ranges within that mapping. The complete byte count is checked
 against the per-worker budget before any platform mutation. The CUDA backend
-then registers every range with the read-only host-registration flag and
-verifies that PyTorch recognizes every mapped tensor as pinned. Failure rolls
-back any ranges already registered and preserves the existing two-slot staging
-path. Partial direct transfer is deliberately unsupported. A rollback failure
-aborts worker initialization so a still-registered range is retained until
-process/context teardown rather than being unsafely unmapped.
+first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every
+range with the read-only host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. A device without that capability
+retains bounded staging; vLLM-Omni does not replace immutable mappings with a
+writable or copy-on-write CUDA-specific layout. Other registration failures
+roll back any ranges already registered and preserve the existing two-slot
+staging path. Partial direct transfer is deliberately unsupported. A rollback
+failure aborts worker initialization so a still-registered range is retained
+until process/context teardown rather than being unsafely unmapped.
 
 On success, the hooks copy each immutable tensor view directly into its offset
 in the existing flattened rotating device buffer on the copy stream. Parameter

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -83,8 +83,10 @@ controls are intentionally not separate CLI flags.
   that tradeoff, or use a validated registered path.
 - A positive registration budget must cover the complete page-aligned mapping
   reported in the worker log. Registration is all-or-nothing.
-- CUDA is the first registration backend. Platforms without an equivalent
-  implementation continue to use bounded staging.
+- CUDA is the first registration backend and requires read-only host
+  registration support for immutable cache mappings. Platforms or devices
+  without that capability continue to use bounded staging; a positive budget
+  requests direct H2D but does not guarantee it.
 - Registration is process-local but does not duplicate the underlying file
   pages. Each worker must still satisfy its platform and OS page-locking limits.
 - Shutdown unregisters host ranges before closing their mmap handles.

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import ctypes
 from types import SimpleNamespace
 
 import pytest
@@ -58,12 +59,23 @@ class _FakeRuntime:
         self,
         register_results: list[int | Exception],
         unregister_results: list[int | Exception] | None = None,
+        *,
+        read_only_registration_supported: bool = True,
+        attribute_result: int = 0,
     ) -> None:
         self._register_results = iter(register_results)
         self._unregister_results = iter(unregister_results or [])
+        self._read_only_registration_supported = read_only_registration_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_registration_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)
@@ -83,6 +95,11 @@ class _FakeRuntime:
         return f"error-{error}"
 
 
+@pytest.fixture(autouse=True)
+def _current_cuda_device(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
+
+
 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)
@@ -126,6 +143,40 @@ def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytes
     assert runtime.registered == []
 
 
+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)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="does not support read-only host registration"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.attribute_queries == [
+        (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
+    ]
+    assert runtime.registered == []
+
+
+def test_registration_consumes_capability_query_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], attribute_result=7)
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    consumed_errors: list[int] = []
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert consumed_errors == [7]
+    assert runtime.registered == []
+
+
 def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0, 7])

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -19,6 +19,7 @@ from .host_registration import (
 )
 
 _CUDA_HOST_REGISTER_READ_ONLY = 0x08
+_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
 
 CudaHostRegistrationError = HostRegistrationError
 CudaHostRegistrationBudgetError = HostRegistrationBudgetError
@@ -133,6 +134,37 @@ def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
     return message
 
 
+def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
+    """Query support required by immutable file-backed cache 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 already-existing file-backed CPU tensors."""
 
@@ -169,6 +201,11 @@ class CudaHostRegistration:
             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 "
+                "host weight cache mappings"
+            )
 
         registered: list[_AddressRange] = []
         try: