diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a5fcbdd22c2e..62faf997cf1b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -512,6 +512,43 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) static nb::object sLogicError = nb::exception(m, "LogicError"); static nb::object sResourceBusyError = nb::exception(m, "ResourceBusyError"); static nb::object sOutOfPagesError = nb::exception(m, "OutOfPagesError"); + static nb::object sCuError = nb::exception(m, "CuError"); + // Default attribute so the class mirrors the pure-Python CuError surface. + sCuError.attr("error_code") = nb::none(); + + // Translate kv::CuError so the Python instance carries the numeric CUDA + // error code (mirrors the pure-Python CuError.error_code). Registered after + // the nb::exception auto-translator so it is tried first. + nb::register_exception_translator( + [](std::exception_ptr const& p, void*) + { + try + { + if (p) + { + std::rethrow_exception(p); + } + } + catch (kv::CuError const& e) + { + nb::object inst = sCuError(nb::str(e.what())); + // Match Python's error_code type (cuda.bindings.driver.CUresult) + // when available; fall back to a plain int otherwise. + nb::object code; + try + { + nb::object cuResult = nb::module_::import_("cuda.bindings.driver").attr("CUresult"); + code = cuResult(static_cast(e.errorCode)); + } + catch (nb::python_error const&) + { + PyErr_Clear(); + code = nb::cast(static_cast(e.errorCode)); + } + inst.attr("error_code") = code; + PyErr_SetObject(sCuError.ptr(), inst.ptr()); + } + }); // Map kv::AssertionError to Python's builtin AssertionError so shared tests see // the same exception type as the pure-Python backend (which uses `assert`). @@ -522,7 +559,9 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) try { if (p) + { std::rethrow_exception(p); + } } catch (kv::AssertionError const& e) { diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 2d5b90468bdf..7b9322b1f63c 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -215,11 +215,11 @@ class _KVCacheManagerConfigFieldSpec: _cpp_introspection = getattr(_cpp, "_introspection", None) _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names) PlannedDropHandle = _cpp.PlannedDropHandle + CuError = _cpp.CuError # Symbols added on main that are not yet ported to the C++ backend. # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) - CuError = getattr(_cpp, "CuError", RuntimeError) OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py index 20c59e903020..7c037cb36d3d 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py @@ -51,6 +51,9 @@ def __init__(self, error_code: drv.CUresult) -> None: err_str = "" super().__init__(f"CUDA driver error: {error_code} ({err_str})") + def __reduce__(self) -> tuple[type["CuError"], tuple[drv.CUresult]]: + return (self.__class__, (self.error_code,)) + class ResourceBusyError(Exception): pass