[Diffusion] Clear handled CUDA registration errors

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

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -130,6 +130,8 @@ def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0, 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(
@@ -141,6 +143,7 @@ def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch
         )
 
     assert runtime.unregistered == [0x1000]
+    assert consumed_errors == [7]
 
 
 def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -164,6 +167,8 @@ def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.Monke
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0, 7], unregister_results=[9])
     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(CudaHostRegistrationCleanupError, match="rollback errors"):
         CudaHostRegistration.create(
@@ -174,6 +179,35 @@ def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.Monke
             max_bytes=8192,
         )
 
+    assert consumed_errors == [7, 9]
+
+
+@pytest.mark.parametrize("pending_error", [0, 801])
+def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
+    monkeypatch: pytest.MonkeyPatch,
+    pending_error: int,
+) -> None:
+    class _GetLastError:
+        def __call__(self) -> int:
+            return pending_error
+
+    runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
+    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
+
+    registration_module._consume_last_cuda_error(801)
+
+
+def test_consume_last_cuda_error_rejects_unrelated_pending_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    class _GetLastError:
+        def __call__(self) -> int:
+            return 700
+
+    runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
+    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="cudaGetLastError reported 700"):
+        registration_module._consume_last_cuda_error(801)
+
 
 def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
@@ -219,6 +253,8 @@ def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatc
     monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
     runtime = _FakeRuntime([0], unregister_results=[9, 0])
     monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    consumed_errors: list[int] = []
+    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
     registration = CudaHostRegistration.create(
         {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
         max_bytes=4096,
@@ -227,3 +263,4 @@ def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatc
     assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
     assert registration.close() == []
     assert runtime.unregistered == [0x1000, 0x1000]
+    assert consumed_errors == [9]

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import ctypes
 import mmap
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
@@ -99,6 +100,39 @@ def _error_message(runtime: _CudaRuntime, error: int) -> str:
     return str(message)
 
 
+def _consume_last_cuda_error(expected_error: int) -> None:
+    """Clear a handled CUDA Runtime error before returning to PyTorch.
+
+    PyTorch's ``torch.cuda.cudart()`` binding exposes host registration but not
+    ``cudaGetLastError``. Resolve the already-loaded process symbol so this call
+    clears the same CUDA Runtime instance that produced the registration error.
+    A different pending error is not safe to hide behind the staging fallback.
+    """
+    try:
+        get_last_error = ctypes.CDLL(None).cudaGetLastError
+        get_last_error.argtypes = []
+        get_last_error.restype = ctypes.c_int
+        pending_error = int(get_last_error())
+    except (AttributeError, OSError) as exc:
+        raise HostRegistrationCleanupError("cannot clear CUDA's pending error after host-registration failure") from exc
+
+    if pending_error not in (0, expected_error):
+        raise HostRegistrationCleanupError(
+            f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
+        )
+
+
+def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
+    """Format and consume one nonzero CUDA Runtime return code."""
+    error_code = int(error)
+    message = _error_message(runtime, error)
+    try:
+        _consume_last_cuda_error(error_code)
+    except HostRegistrationCleanupError as exc:
+        raise HostRegistrationCleanupError(f"{message}; {exc}") from exc
+    return message
+
+
 class CudaHostRegistration:
     """Own CUDA registrations for already-existing file-backed CPU tensors."""
 
@@ -147,7 +181,7 @@ class CudaHostRegistration:
                 if int(error) != 0:
                     raise HostRegistrationError(
                         "cudaHostRegister(read-only) failed for "
-                        f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}"
+                        f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
                     )
                 registered.append(region)
 
@@ -167,7 +201,7 @@ class CudaHostRegistration:
                     error = runtime.cudaHostUnregister(region.start)
                     if int(error) != 0:
                         rollback_errors.append(
-                            f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}"
+                            f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
                         )
                 except Exception as rollback_exc:
                     rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
@@ -200,7 +234,7 @@ class CudaHostRegistration:
                 error = self._runtime.cudaHostUnregister(region.start)
                 if int(error) != 0:
                     errors.append(
-                        f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}"
+                        f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
                     )
                     failed.append(region)
             except Exception as exc: