[Python] Reduce latency of NIXL import

src/api/python/_api.py

@@ -13,19 +13,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import pickle
+import sys
 from enum import Enum
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
-import torch
 
 from . import _bindings as nixlBind  # type: ignore
 from .logging import get_logger
 
+if TYPE_CHECKING:
+    import torch
+
 # Get logger using centralized configuration
 logger = get_logger(__name__)
 
+
+def _is_torch_tensor(obj: object) -> bool:
+    """Return True if obj is a torch.Tensor and torch is already imported.
+
+    Importing torch costs ~1s and dominates `import nixl`. A value can only be
+    a torch.Tensor if torch has already been imported elsewhere, so we consult
+    sys.modules instead of importing torch here.
+    """
+    torch = sys.modules.get("torch")
+    return torch is not None and isinstance(obj, torch.Tensor)
+
+
 DEFAULT_COMM_PORT = nixlBind.DEFAULT_COMM_PORT
 
 
@@ -983,7 +1000,7 @@ def get_xfer_descs(
                     "Nx3 shape required for transfer descriptor list from numpy array"
                 )
                 new_descs = None
-        elif isinstance(descs, torch.Tensor):
+        elif _is_torch_tensor(descs):
             if descs.is_contiguous():
                 mem_type = self._tensor_mem_type(descs)
                 base_addr = descs.data_ptr()
@@ -997,7 +1014,7 @@ def get_xfer_descs(
             else:
                 logger.error("Please use a list of contiguous Tensors")
                 new_descs = None
-        elif isinstance(descs[0], torch.Tensor):  # List[torch.Tensor]:
+        elif _is_torch_tensor(descs[0]):  # List[torch.Tensor]:
             tensor_type = descs[0].device
             dlist = np.zeros((len(descs), 3), dtype=np.uint64)
 
@@ -1066,7 +1083,7 @@ def get_reg_descs(
                     "Nx3 shape required for transfer descriptor list from numpy array"
                 )
                 new_descs = None
-        elif isinstance(descs, torch.Tensor):
+        elif _is_torch_tensor(descs):
             if descs.is_contiguous():
                 mem_type = self._tensor_mem_type(descs)
                 base_addr = descs.data_ptr()
@@ -1080,7 +1097,7 @@ def get_reg_descs(
             else:
                 logger.error("Please use a list of contiguous Tensors")
                 new_descs = None
-        elif isinstance(descs[0], torch.Tensor):  # List[torch.Tensor]:
+        elif _is_torch_tensor(descs[0]):  # List[torch.Tensor]:
             tensor_type = descs[0].device
             dlist = np.zeros((len(descs), 3), dtype=np.uint64)
 

src/bindings/python/nixl-meta/meson.build

@@ -29,6 +29,7 @@ pyproject_toml = configure_file(
 )
 
 readme_md = fs.copyfile('README.md')
+nixl_meta_utils_copy = fs.copyfile('nixl_meta_utils.py')
 
 source_root = meson.project_source_root()
 root_license_path = join_paths(source_root, 'LICENSE')
@@ -50,6 +51,7 @@ if uv.found()
       pyproject_toml,
       readme_md,
       license_path,
+      nixl_meta_utils_copy,
       nixl_init_copy,
       nixl_api_copy,
       nixl_logging_copy,

src/bindings/python/nixl-meta/nixl/__init__.py

@@ -17,16 +17,11 @@
 import sys
 from typing import TYPE_CHECKING
 
-
-def _get_torch_cuda_major() -> int | None:
-    """Return the CUDA major version that torch was built for, or None."""
-    from torch.version import cuda as _torch_cuda_ver
-
-    return int(_torch_cuda_ver.split(".")[0]) if _torch_cuda_ver else None
+from nixl_meta_utils import detect_cuda_major
 
 
 def _load_cuda_backend() -> str:
-    cuda_major = _get_torch_cuda_major()
+    cuda_major = detect_cuda_major()
     if cuda_major is not None:
         pip_name = f"nixl-cu{cuda_major}"
         mod_name = f"nixl_cu{cuda_major}"
@@ -36,9 +31,9 @@ def _load_cuda_backend() -> str:
             if e.name != mod_name:
                 raise
             raise ImportError(
-                f"torch reports CUDA {cuda_major} but {pip_name} is not installed"
+                f"detected CUDA {cuda_major} but {pip_name} is not installed"
             ) from e
-    # CPU-only torch — use whatever backend is installed
+    # No CUDA stack detected — use whatever backend is installed.
     for mod_name in ("nixl_cu13", "nixl_cu12"):
         try:
             return importlib.import_module(mod_name).__name__

src/bindings/python/nixl-meta/nixl_ep/__init__.py

@@ -19,16 +19,11 @@
 import sys
 from typing import TYPE_CHECKING
 
-
-def _get_torch_cuda_major() -> int | None:
-    """Return the CUDA major version that torch was built for, or None."""
-    from torch.version import cuda as _torch_cuda_ver
-
-    return int(_torch_cuda_ver.split(".")[0]) if _torch_cuda_ver else None
+from nixl_meta_utils import detect_cuda_major
 
 
 def _load_ep_module() -> str:
-    cuda_major = _get_torch_cuda_major()
+    cuda_major = detect_cuda_major()
     if cuda_major is not None:
         pip_name = f"nixl-cu{cuda_major}"
         mod_name = f"nixl_ep_cu{cuda_major}"
@@ -38,9 +33,9 @@ def _load_ep_module() -> str:
             if e.name != mod_name:
                 raise
             raise ImportError(
-                f"torch reports CUDA {cuda_major} but {pip_name} is not installed"
+                f"detected CUDA {cuda_major} but {pip_name} is not installed"
             ) from e
-    # CPU-only torch — use whatever backend is installed
+    # No CUDA stack detected — use whatever backend is installed.
     errors: list[BaseException] = []
     for mod_name in ("nixl_ep_cu13", "nixl_ep_cu12"):
         try:

src/bindings/python/nixl-meta/nixl_meta_utils.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared helpers for the NIXL meta packages (``nixl`` and ``nixl_ep``)."""
+
+import importlib.util
+import pathlib
+import sys
+
+
+def detect_cuda_major() -> int | None:
+    """CUDA major version used to select the nixl_cuXX backend wheel.
+
+    Read from whichever CUDA stack the process has already imported, so that
+    importing a NIXL meta package does not itself pull in torch (a ~1s cost).
+    If nothing CUDA is loaded yet, read torch's build version off disk (still
+    torch-free); only if that fails do we import torch and use its official API
+    as a last resort.
+    """
+
+    def major(version: str | None) -> int | None:
+        return int(version.split(".")[0]) if version else None
+
+    # torch already imported: use its official API.
+    torch = sys.modules.get("torch")
+    if torch is not None:
+        version = major(getattr(getattr(torch, "version", None), "cuda", None))
+        if version is not None:
+            return version
+
+    # cuda-python already imported: use it.
+    cuda_bindings = sys.modules.get("cuda.bindings")
+    if cuda_bindings is not None:
+        version = major(getattr(cuda_bindings, "__version__", None))
+        if version is not None:
+            return version
+
+    # cupy already imported: use it.
+    cupy = sys.modules.get("cupy")
+    if cupy is not None:
+        try:
+            return cupy.cuda.runtime.runtimeGetVersion() // 1000
+        except Exception:
+            pass
+
+    # torch installed but not imported: read build version off disk.
+    version = major(_torch_cuda_version_from_disk())
+    if version is not None:
+        return version
+
+    # Last resort: import torch and use its official API. Slow.
+    try:
+        from torch.version import cuda as torch_cuda
+
+        return major(torch_cuda)
+    except ImportError:
+        return None
+
+
+def _torch_cuda_version_from_disk() -> str | None:
+    """Return torch's build CUDA version (e.g. "12.6") without full torch import.
+
+    ``find_spec`` locates the package without running its ``__init__``, and the
+    standalone module name keeps importlib from importing the torch package; we
+    then exec only the tiny ``version.py`` to read its ``cuda`` attribute.
+    """
+    try:
+        spec = importlib.util.find_spec("torch")
+    except (ImportError, ValueError):
+        return None
+    if spec is None or not spec.origin:
+        return None
+
+    version_py = pathlib.Path(spec.origin).parent / "version.py"
+    try:
+        version_spec = importlib.util.spec_from_file_location(
+            "_nixl_torch_version", version_py
+        )
+        if version_spec is None or version_spec.loader is None:
+            return None
+        module = importlib.util.module_from_spec(version_spec)
+        version_spec.loader.exec_module(module)
+    except Exception:
+        return None
+    return getattr(module, "cuda", None)

src/bindings/python/nixl-meta/pyproject.toml.in

@@ -37,3 +37,4 @@ cu13 = ["nixl-cu13==@VERSION@"]
 
 [tool.setuptools]
 packages = ["nixl", "nixl_ep"]
+py-modules = ["nixl_meta_utils"]