[NVIDIA] Add opt-in NVML-based CUDA availability check

CudaDriver.is_active() calls cuInit(), which initializes CUDA driver state in a parent process and can prevent a subsequently forked child from initializing CUDA.

Add an opt-in NVML probe so callers that need to fork can check NVIDIA device availability without initializing CUDA when NVML returns a conclusive result. Keep the existing CUDA probe as the default and as the fallback for inconclusive NVML checks.

Author: warrendeng

README.md

@@ -246,6 +246,8 @@ See [`python/triton/knobs.py`](python/triton/knobs.py) for the full list of conf
 - `LLVM_EXTRACT_DI_LOCAL_VARIABLES` emit full debug info, allowing for eval of values in gpu debuggers (ie cuda-gdb, rocm-gdb etc)
 - `TRITON_DEFAULT_BACKEND=<backend>` optionally sets the default backend used by Triton when
   constructing the active driver (i.e., `triton.runtime.driver.active`).
+- `TRITON_NVML_BASED_CUDA_CHECK=1` makes NVIDIA driver availability checks use NVML when possible
+  to avoid initializing CUDA before a process forks.
 
 > [!NOTE]
 > Some of these environment variables don't have a knob in `knobs.py`-- those are only relevant to the C++ layer(s), hence they don't exist in the python layer.

python/test/unit/runtime/test_driver.py

@@ -5,6 +5,7 @@
 
 import triton
 import triton.language as tl
+import triton.backends.nvidia.driver as cuda_driver
 from triton.backends.driver import GPUDriver, expand_signature, wrap_handle_tensordesc_impl
 
 
@@ -19,6 +20,18 @@ def test_is_lazy():
     utils = triton.runtime.driver.active.utils  # noqa: F841
 
 
+def test_cuda_driver_is_active_with_nvml(monkeypatch):
+    cuda_calls = []
+    monkeypatch.setattr(cuda_driver, "_device_count_nvml", lambda: 1)
+    monkeypatch.setattr(cuda_driver, "_cuda_driver_is_active", lambda: cuda_calls.append(None) or False)
+
+    with triton.knobs.nvidia.scope():
+        triton.knobs.nvidia.use_nvml_cuda_check = True
+        assert cuda_driver.CudaDriver.is_active() is True
+
+    assert cuda_calls == []
+
+
 def test_profile_scratch_stream_zero_uses_default_stream(monkeypatch):
 
     class Scratch:

python/triton/knobs.py

@@ -511,6 +511,7 @@ class nvidia_knobs(base_knobs):
 
     libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH")
     libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH")
+    use_nvml_cuda_check: env_bool = env_bool("TRITON_NVML_BASED_CUDA_CHECK")
 
 
 class amd_knobs(base_knobs):

third_party/nvidia/backend/driver.py

@@ -51,6 +51,115 @@ def library_dirs():
     return [libdevice_dir, *libcuda_dirs()]
 
 
+# Adapted from PyTorch's NVML-based CUDA availability check:
+# https://github.com/pytorch/pytorch/blob/70d99e998b4955e0049d13a98d77ae1b14db1f45/torch/cuda/__init__.py
+def _parse_visible_devices():
+    var = os.getenv("CUDA_VISIBLE_DEVICES")
+    if var is None:
+        return list(range(64))
+
+    def _strtoul(s):
+        if not s:
+            return -1
+        for idx, c in enumerate(s):
+            if not (c.isdigit() or (idx == 0 and c in "+-")):
+                break
+            if idx + 1 == len(s):
+                idx += 1
+        return int(s[:idx]) if idx > 0 else -1
+
+    def parse_list_with_prefix(prefix):
+        rc = []
+        for elem in var.split(","):
+            if elem in rc:
+                return []
+            if not elem.startswith(prefix):
+                break
+            rc.append(elem)
+        return rc
+
+    if var.startswith("GPU-"):
+        return parse_list_with_prefix("GPU-")
+    if var.startswith("MIG-"):
+        return parse_list_with_prefix("MIG-")
+
+    rc = []
+    for elem in var.split(","):
+        x = _strtoul(elem.strip())
+        if x in rc:
+            return []
+        if x < 0:
+            break
+        rc.append(x)
+    return rc
+
+
+def _raw_device_count_nvml():
+    nvml_h = ctypes.CDLL("libnvidia-ml.so.1")
+    if nvml_h.nvmlInit() != 0:
+        return -1
+    count = ctypes.c_int(-1)
+    if nvml_h.nvmlDeviceGetCount_v2(ctypes.byref(count)) != 0:
+        return -1
+    return count.value
+
+
+def _raw_device_uuid_nvml():
+    nvml_h = ctypes.CDLL("libnvidia-ml.so.1")
+    if nvml_h.nvmlInit() != 0:
+        return None
+    count = ctypes.c_int(-1)
+    if nvml_h.nvmlDeviceGetCount_v2(ctypes.byref(count)) != 0:
+        return None
+
+    uuids = []
+    for idx in range(count.value):
+        device = ctypes.c_void_p()
+        if nvml_h.nvmlDeviceGetHandleByIndex_v2(idx, ctypes.byref(device)) != 0:
+            return None
+        buf = ctypes.create_string_buffer(96)
+        if nvml_h.nvmlDeviceGetUUID(device, buf, len(buf)) != 0:
+            return None
+        uuids.append(buf.raw.decode("ascii").strip("\0"))
+    return uuids
+
+
+def _transform_uuid_to_ordinals(candidates, uuids):
+    ordinals = []
+    for candidate in candidates:
+        matches = [idx for idx, uuid in enumerate(uuids) if uuid.startswith(candidate)]
+        if len(matches) != 1:
+            break
+        if matches[0] in ordinals:
+            return []
+        ordinals.append(matches[0])
+    return ordinals
+
+
+def _device_count_nvml():
+    visible_devices = _parse_visible_devices()
+    if not visible_devices:
+        return 0
+    try:
+        if isinstance(visible_devices[0], str):
+            if visible_devices[0].startswith("MIG-"):
+                return -1
+            uuids = _raw_device_uuid_nvml()
+            if uuids is None:
+                return -1
+            visible_devices = _transform_uuid_to_ordinals(visible_devices, uuids)
+        else:
+            raw_count = _raw_device_count_nvml()
+            if raw_count <= 0:
+                return raw_count
+            for idx, value in enumerate(visible_devices):
+                if value >= raw_count:
+                    return idx
+    except (OSError, AttributeError):
+        return -1
+    return len(visible_devices)
+
+
 def _cuda_driver_is_active():
     candidates = ["libcuda.so.1"]
     try:
@@ -381,6 +490,10 @@ def get_device_interface(self):
 
     @staticmethod
     def is_active():
+        if knobs.nvidia.use_nvml_cuda_check:
+            count = _device_count_nvml()
+            if count >= 0:
+                return count > 0
         return _cuda_driver_is_active()
 
     def map_python_to_cpp_type(self, ty: str) -> str: