Polygraphy: find the CUDA runtime from CUDA 13+ pip wheels

tools/Polygraphy/CHANGELOG.md

@@ -3,6 +3,13 @@
 Dates are in YYYY-MM-DD format.
 
 
+## Unreleased
+### Fixed
+- Polygraphy now finds the CUDA runtime installed by the CUDA 13+ `nvidia-cuda-runtime` pip wheel, which
+  places it under `nvidia/cu13` instead of `nvidia/cuda_runtime`. Previously, without a CUDA Toolkit or
+  `LD_LIBRARY_PATH` entry, TensorRT runners failed with `OSError: libcudart.so: cannot open shared object file`.
+
+
 ## v0.53.6 (2026-09-22)
 ### Added
 - `Comparator.run()` now accepts a `save_input_blob_path` parameter, exposed on the CLI as

tools/Polygraphy/polygraphy/cuda/cuda.py

@@ -16,6 +16,7 @@
 #
 import contextlib
 import ctypes
+import glob
 import importlib.util
 import os
 import platform
@@ -32,24 +33,27 @@ def void_ptr(val=None):
     return ctypes.c_void_p(val)
 
 
-def _cuda_runtime_wheel_dirs(subdir):
-    # The pip `nvidia-cuda-runtime-cuXX` wheel installs the runtime under
-    # `<site-packages>/nvidia/cuda_runtime/{bin (Windows), lib (Linux)}`.
-    dirs = [
-        # Some `sys.path` entries are files (a console-script `.exe`, a zipapp), not dirs; skip those.
-        os.path.join(p, "nvidia", "cuda_runtime", subdir)
-        for p in sys.path
-        if p and os.path.isdir(p)
+def _cuda_runtime_wheel_dirs(*subdirs):
+    # The pip CUDA runtime wheel installs the runtime under `<site-packages>/nvidia/cuda_runtime`
+    # for CUDA 12 (`nvidia-cuda-runtime-cu12`) and `<site-packages>/nvidia/cu<major>` for CUDA 13+
+    # (`nvidia-cuda-runtime`).
+    # Some `sys.path` entries are files (a console-script `.exe`, a zipapp), not dirs; skip those.
+    nvidia_dirs = [
+        os.path.join(p, "nvidia") for p in sys.path if p and os.path.isdir(p)
     ]
+    pkg_dirs = [os.path.join(d, "cuda_runtime") for d in nvidia_dirs]
     # find_spec also locates the package when it is import-reachable (e.g. via a `.pth`) but its
     # directory is not itself on `sys.path`. Best-effort: an uninstalled wheel just raises.
     with contextlib.suppress(Exception):
         spec = importlib.util.find_spec("nvidia.cuda_runtime")
         if spec is not None and spec.submodule_search_locations:
-            dirs += [
-                os.path.join(loc, subdir) for loc in spec.submodule_search_locations
-            ]
-    return dirs
+            pkg_dirs += list(spec.submodule_search_locations)
+    pkg_dirs += [
+        pkg_dir
+        for d in nvidia_dirs
+        for pkg_dir in sorted(glob.glob(os.path.join(d, "cu[0-9]*")))
+    ]
+    return [os.path.join(d, subdir) for d in pkg_dirs for subdir in subdirs]
 
 
 def _find_cuda_lib_dirs():
@@ -79,7 +83,7 @@ def _find_cuda_lib_dirs():
             if root
             for path in [root, *[os.path.join(root, s) for s in bin_subdirs]]
         ]
-        cuda_paths += _cuda_runtime_wheel_dirs("bin")
+        cuda_paths += _cuda_runtime_wheel_dirs("bin", os.path.join("bin", "x86_64"))
         cuda_paths += os.environ.get("PATH", "").split(os.path.pathsep)
     else:
         cuda_paths = [
@@ -139,13 +143,15 @@ def __init__(self):
                 hint = (
                     "Ensure the CUDA Toolkit is installed and that either `CUDA_PATH` points to its "
                     "root or the directory containing `cudart64_*.dll` is on `PATH`. "
-                    "Alternatively, `pip install nvidia-cuda-runtime-cuXX` provides the runtime."
+                    "Alternatively, `pip install nvidia-cuda-runtime` (CUDA 13+) or "
+                    "`nvidia-cuda-runtime-cu12` provides the runtime."
                 )
             else:
                 hint = (
                     "Ensure the CUDA Toolkit is installed and that the directory containing "
                     "`libcudart.so*` is on `LD_LIBRARY_PATH`. "
-                    "Alternatively, `pip install nvidia-cuda-runtime-cuXX` provides the runtime."
+                    "Alternatively, `pip install nvidia-cuda-runtime` (CUDA 13+) or "
+                    "`nvidia-cuda-runtime-cu12` provides the runtime."
                 )
             log_func(
                 f"Could not find the CUDA runtime library.\n{hint}\n"

tools/Polygraphy/tests/cuda/test_cuda.py

@@ -127,6 +127,28 @@ class FakeSpec:
         dirs = _find_cuda_lib_dirs()
         assert os.path.join(pkg_dir, "bin") in dirs
 
+    @pytest.mark.parametrize(
+        "platform_name,lib_subdir",
+        [("linux", "lib"), ("win32", os.path.join("bin", "x86_64"))],
+    )
+    def test_cuda13_runtime_wheel_dirs_included(
+        self, monkeypatch, tmp_path, platform_name, lib_subdir
+    ):
+        # CUDA 13+ runtime wheels install under `nvidia/cu<major>` instead of `nvidia/cuda_runtime`;
+        # other `nvidia/cu*` packages (e.g. cuBLAS) must not be picked up.
+        (tmp_path / "nvidia" / "cu13").mkdir(parents=True)
+        (tmp_path / "nvidia" / "cublas").mkdir()
+        monkeypatch.setattr("sys.platform", platform_name)
+        monkeypatch.setattr("sys.path", [str(tmp_path)])
+        monkeypatch.setattr("platform.machine", lambda: "AMD64")
+        monkeypatch.setattr(importlib.util, "find_spec", lambda name: None)
+        monkeypatch.setattr(os, "environ", {})
+
+        dirs = _find_cuda_lib_dirs()
+
+        assert os.path.join(str(tmp_path), "nvidia", "cu13", lib_subdir) in dirs
+        assert not any("cublas" in d for d in dirs)
+
     def test_windows_deduplicates_preserving_order(self, monkeypatch):
         # CUDA_PATH and a versioned variable can point at the same root; the candidate list must
         # not contain duplicates, and first-seen order is preserved.