cuda.core: raise CUDAOutOfMemoryError, a MemoryError, on out-of-memory

cuda_core/cuda/core/__init__.py

@@ -168,10 +168,11 @@ class _PatchedProperty(metaclass=_PatchedPropMeta):
 from cuda.core._stream import __all__ as _stream_all
 from cuda.core._tensor_map import *
 from cuda.core._tensor_map import __all__ as _tensor_map_all
-from cuda.core._utils.cuda_utils import CUDAError, CUDAWarning, NVRTCError
+from cuda.core._utils.cuda_utils import CUDAError, CUDAOutOfMemoryError, CUDAWarning, NVRTCError
 
 __all__ = [
     "CUDAError",
+    "CUDAOutOfMemoryError",
     "CUDAWarning",
     "NVRTCError",
     *_context_all,

cuda_core/cuda/core/_utils/cuda_utils.pyi

@@ -22,6 +22,17 @@ class CUDAError(Exception):
     documentation for the guarantees an exception provides.
     """
 
+class CUDAOutOfMemoryError(CUDAError, MemoryError):
+    """Raised when a CUDA driver or runtime call fails for lack of memory.
+
+    ``CUDA_ERROR_OUT_OF_MEMORY`` and ``cudaErrorMemoryAllocation`` raise this
+    class. It is both a :class:`CUDAError` and a :class:`MemoryError`, so code
+    that catches :class:`MemoryError` to retry with a smaller request handles
+    ``cuda.core`` the same way as other array libraries.
+
+    .. versionadded:: 1.3.0
+    """
+
 class CUDAWarning(RuntimeWarning):
     """Warning issued when ``cuda.core`` hits a CUDA error it cannot raise.
 

cuda_core/cuda/core/_utils/cuda_utils.pyx

@@ -45,6 +45,18 @@ class CUDAError(Exception):
     """
 
 
+class CUDAOutOfMemoryError(CUDAError, MemoryError):
+    """Raised when a CUDA driver or runtime call fails for lack of memory.
+
+    ``CUDA_ERROR_OUT_OF_MEMORY`` and ``cudaErrorMemoryAllocation`` raise this
+    class. It is both a :class:`CUDAError` and a :class:`MemoryError`, so code
+    that catches :class:`MemoryError` to retry with a smaller request handles
+    ``cuda.core`` the same way as other array libraries.
+
+    .. versionadded:: 1.3.0
+    """
+
+
 class CUDAWarning(RuntimeWarning):
     """Warning issued when ``cuda.core`` hits a CUDA error it cannot raise.
 
@@ -163,6 +175,7 @@ cdef int _raise_nvjitlink_error(
 
 
 cdef object _RUNTIME_SUCCESS = runtime.cudaError_t.cudaSuccess
+cdef object _RUNTIME_OUT_OF_MEMORY = runtime.cudaError_t.cudaErrorMemoryAllocation
 cdef object _NVRTC_SUCCESS = nvrtc.nvrtcResult.NVRTC_SUCCESS
 
 
@@ -184,14 +197,18 @@ cpdef int _check_driver_error(cydriver.CUresult error) except?-1 nogil:
         raise CUDAError(f"UNEXPECTED ERROR CODE: {error}")
     desc_err = cydriver.cuGetErrorString(error, &desc)
     with gil:
+        exc_type = (
+            CUDAOutOfMemoryError if error == cydriver.CUresult.CUDA_ERROR_OUT_OF_MEMORY
+            else CUDAError
+        )
         # TODO: consider lower this to Cython
         expl = DRIVER_CU_RESULT_EXPLANATIONS.get(int(error))
         if expl is not None:
-            exc = CUDAError(f"{name.decode()}: {expl}")
+            exc = exc_type(f"{name.decode()}: {expl}")
         elif desc_err != cydriver.CUresult.CUDA_SUCCESS:
-            exc = CUDAError(name.decode())
+            exc = exc_type(name.decode())
         else:
-            exc = CUDAError(f"{name.decode()}: {desc.decode()}")
+            exc = exc_type(f"{name.decode()}: {desc.decode()}")
         if detail != NULL:
             _attach_detail(exc, detail.decode())
         raise exc
@@ -203,14 +220,15 @@ cpdef int _check_runtime_error(error) except?-1:
     # `_check_error()` reaches this path only for `runtime.cudaError_t` values.
     # Use the enum name directly because Windows hybrid cudart can lag that table.
     name = error.name
+    exc_type = CUDAOutOfMemoryError if error == _RUNTIME_OUT_OF_MEMORY else CUDAError
     expl = RUNTIME_CUDA_ERROR_EXPLANATIONS.get(int(error))
     if expl is not None:
-        raise CUDAError(f"{name}: {expl}")
+        raise exc_type(f"{name}: {expl}")
     desc_err, desc = runtime.cudaGetErrorString(error)
     if desc_err != _RUNTIME_SUCCESS:
-        raise CUDAError(f"{name}")
+        raise exc_type(f"{name}")
     desc = desc.decode()
-    raise CUDAError(f"{name}: {desc}")
+    raise exc_type(f"{name}: {desc}")
 
 
 cpdef int _check_nvrtc_error(error, handle=None) except?-1:

cuda_core/docs/source/api.rst

@@ -326,6 +326,7 @@ reported as a warning instead.
    :toctree: generated/
 
    CUDAError
+   CUDAOutOfMemoryError
    NVRTCError
    CUDAWarning
 

cuda_core/docs/source/error_handling.rst

@@ -18,8 +18,10 @@ A CUDA driver, runtime, NVRTC, NVVM or nvJitLink call that fails raises an
 exception whose message contains the error name and its description:
 :class:`CUDAError` for driver and runtime failures, its subclass
 :class:`NVRTCError` for NVRTC failures, and the ``cuda.bindings`` error types
-for NVVM and nvJitLink failures. Both ``cuda.core`` classes are importable from
-the top-level ``cuda.core`` namespace. Invalid arguments and misuse raise the
+for NVVM and nvJitLink failures. A driver or runtime call that runs out of
+memory raises :class:`CUDAOutOfMemoryError`, which derives from both
+:class:`CUDAError` and :class:`MemoryError`. All ``cuda.core`` classes are
+importable from the top-level ``cuda.core`` namespace. Invalid arguments and misuse raise the
 usual Python exception types (``TypeError``, ``ValueError``, ``RuntimeError``).
 
 When a ``cuda.core`` call raises, the following hold:

cuda_core/docs/source/release/1.3.0-notes.rst

@@ -76,6 +76,12 @@ New features
   top-level ``cuda.core`` namespace, alongside :class:`CUDAWarning`. Previously
   they were only available from a private module.
 
+- Added :class:`CUDAOutOfMemoryError`, raised for ``CUDA_ERROR_OUT_OF_MEMORY``
+  and ``cudaErrorMemoryAllocation``. It derives from both :class:`CUDAError`
+  and :class:`MemoryError`, so ``except MemoryError`` catches an allocation
+  that fails in ``cuda.core`` as it does in CuPy, PyTorch and NumPy. Code that
+  catches :class:`CUDAError` is unaffected.
+
 - Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that
   :class:`VirtualMemoryResource` returns. Closing it on a capturing stream
   other than a default stream raises, because virtual memory deallocation

cuda_core/tests/test_cuda_utils.py

@@ -50,6 +50,40 @@ def test_check_runtime_error():
             assert error.name in msg
 
 
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_check_driver_error_out_of_memory_is_memory_error():
+    with pytest.raises(cuda_utils.CUDAOutOfMemoryError) as e:
+        cuda_utils._check_driver_error(driver.CUresult.CUDA_ERROR_OUT_OF_MEMORY)
+    assert isinstance(e.value, cuda_utils.CUDAError)
+    assert isinstance(e.value, MemoryError)
+    assert "CUDA_ERROR_OUT_OF_MEMORY" in str(e.value)
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_check_runtime_error_out_of_memory_is_memory_error():
+    with pytest.raises(cuda_utils.CUDAOutOfMemoryError) as e:
+        cuda_utils._check_runtime_error(runtime.cudaError_t.cudaErrorMemoryAllocation)
+    assert isinstance(e.value, cuda_utils.CUDAError)
+    assert isinstance(e.value, MemoryError)
+    assert "cudaErrorMemoryAllocation" in str(e.value)
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_only_out_of_memory_errors_are_memory_errors():
+    for error in driver.CUresult:
+        if error in (driver.CUresult.CUDA_SUCCESS, driver.CUresult.CUDA_ERROR_OUT_OF_MEMORY):
+            continue
+        with pytest.raises(cuda_utils.CUDAError) as e:
+            cuda_utils._check_driver_error(error)
+        assert not isinstance(e.value, MemoryError), error
+    for error in runtime.cudaError_t:
+        if error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorMemoryAllocation):
+            continue
+        with pytest.raises(cuda_utils.CUDAError) as e:
+            cuda_utils._check_runtime_error(error)
+        assert not isinstance(e.value, MemoryError), error
+
+
 def test_driver_error_enum_has_non_empty_docstring():
 
     doc = driver.CUresult.CUDA_ERROR_INVALID_VALUE.__doc__

cuda_core/tests/test_memory.py

@@ -44,6 +44,7 @@
 
 from cuda.core import (
     Buffer,
+    CUDAOutOfMemoryError,
     CUDAWarning,
     Device,
     DeviceMemoryResource,
@@ -2223,6 +2224,30 @@ def test_vmm_host_location_grow_and_close(init_cuda):
         assert moved.size == 0
 
 
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_device_memory_resource_out_of_memory_is_memory_error(init_cuda):
+    device = Device()
+    if not device.properties.memory_pools_supported:
+        pytest.skip("Device does not support mempool operations")
+    device.set_current()
+
+    mr = DeviceMemoryResource(device, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE))
+    try:
+        # Larger than any device, and than the pool's max_size, which the
+        # driver rounds up to its allocation granularity.
+        with pytest.raises(MemoryError) as e:
+            mr.allocate(1 << 46, stream=device.default_stream)
+        assert isinstance(e.value, CUDAOutOfMemoryError)
+        assert isinstance(e.value, CUDAError)
+        assert "CUDA_ERROR_OUT_OF_MEMORY" in str(e.value)
+
+        # The failed request leaves the pool usable.
+        buffer = mr.allocate(POOL_SIZE // 2, stream=device.default_stream)
+        buffer.close()
+    finally:
+        mr.close()
+
+
 def test_device_memory_resource_with_options(init_cuda):
     device = Device()
     if not device.properties.memory_pools_supported: