From 859003fb97a4de630c8da7b3c3307eae1a098b1b Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 18 Aug 2026 14:48:39 +0800 Subject: [PATCH] [Diffusion] Gate CUDA read-only host registration Signed-off-by: Hongsheng Liu --- .../distributed_layerwise_offload.md | 15 +++--- .../distributed_layerwise_offload.md | 6 ++- .../offloader/test_cuda_host_registration.py | 51 +++++++++++++++++++ .../offloader/cuda_host_registration.py | 37 ++++++++++++++ 4 files changed, 101 insertions(+), 8 deletions(-) diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index d7c47e22a..94684cd81 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/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 diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 2aad95f00..1c27509ad 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/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. diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index 262ce920c..3510de5cc 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/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]) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index 6ee54b02d..1183c9a80 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/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: