Use dlopen to load NCCL EP

build_tools/build_ext.py

@@ -6,6 +6,7 @@
 
 import copy
 import os
+import shutil
 import subprocess
 import sys
 import sysconfig
@@ -171,6 +172,17 @@ def run(self) -> None:
                     self.copy_file(ext, target_dir)
                     os.remove(ext)
 
+                nccl_ep_dir = Path(self.build_lib) / "nccl_ep"
+                if nccl_ep_dir.is_dir():
+                    target_nccl_ep_dir = target_dir / "nccl_ep"
+                    if target_nccl_ep_dir.exists():
+                        shutil.rmtree(target_nccl_ep_dir)
+                    shutil.copytree(
+                        nccl_ep_dir,
+                        target_nccl_ep_dir,
+                    )
+                    shutil.rmtree(nccl_ep_dir)
+
         def build_extensions(self):
             # For core lib + JAX install, fix build_ext from pybind11.setup_helpers
             # to handle CUDA files correctly.

qa/L0_jax_wheel/test.sh

@@ -30,6 +30,12 @@ WHL_BASE="transformer_engine-${VERSION}"
 # Core wheel.
 NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel"
 wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl"
+if python3 -c "from build_tools.utils import nccl_ep_enabled; raise SystemExit(not nccl_ep_enabled())"; then
+    test -f "${WHL_BASE}/transformer_engine/wheel_lib/libnccl_ep.so" || error_exit "Core wheel is missing libnccl_ep.so"
+    test -f "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep.h" || error_exit "Core wheel is missing nccl_ep.h"
+    test -d "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep" || error_exit "Core wheel is missing NCCL EP JIT headers"
+    python3 "$TE_PATH/qa/check_nccl_ep_headers.py" "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include" || error_exit "Core wheel has incomplete NCCL EP JIT headers"
+fi
 sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
 sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
 mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info"

qa/L0_pytorch_wheel/test.sh

@@ -29,6 +29,12 @@ WHL_BASE="transformer_engine-${VERSION}"
 # Core wheel.
 NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel"
 python3 -m wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl"
+if python3 -c "from build_tools.utils import nccl_ep_enabled; raise SystemExit(not nccl_ep_enabled())"; then
+    test -f "${WHL_BASE}/transformer_engine/wheel_lib/libnccl_ep.so" || error_exit "Core wheel is missing libnccl_ep.so"
+    test -f "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep.h" || error_exit "Core wheel is missing nccl_ep.h"
+    test -d "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include/nccl_ep" || error_exit "Core wheel is missing NCCL EP JIT headers"
+    python3 "$TE_PATH/qa/check_nccl_ep_headers.py" "${WHL_BASE}/transformer_engine/wheel_lib/nccl_ep/include" || error_exit "Core wheel has incomplete NCCL EP JIT headers"
+fi
 sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
 sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
 mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info"

qa/check_nccl_ep_headers.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Verify that packaged NCCL EP JIT headers contain all local quoted includes."""
+
+import re
+import sys
+from pathlib import Path
+
+
+_INCLUDE_PATTERN = re.compile(r'^\s*#\s*include\s*"([^"]+)"')
+_EXTERNAL_HEADERS = {"nccl.h", "nccl_device.h"}
+_HEADER_SUFFIXES = {".cuh", ".h", ".hh", ".hpp", ".inc", ".inl"}
+
+
+def main() -> None:
+    include_root = Path(sys.argv[1])
+    jit_root = include_root / "nccl_ep"
+    public_header = include_root / "nccl_ep.h"
+    if not public_header.is_file():
+        raise RuntimeError(f"Missing NCCL EP public header: {public_header}")
+    if not jit_root.is_dir():
+        raise RuntimeError(f"Missing NCCL EP JIT header directory: {jit_root}")
+
+    failures = []
+
+    headers = [public_header, *jit_root.rglob("*")]
+    for header in headers:
+        if not header.is_file() or header.suffix not in _HEADER_SUFFIXES:
+            continue
+        for line_number, line in enumerate(header.read_text().splitlines(), 1):
+            match = _INCLUDE_PATTERN.match(line)
+            if match is None:
+                continue
+            include = match.group(1)
+            if include in _EXTERNAL_HEADERS:
+                continue
+            candidates = (
+                header.parent / include,
+                include_root / include,
+                jit_root / include,
+                jit_root / "device" / include,
+            )
+            if not any(candidate.is_file() for candidate in candidates):
+                failures.append(f"{header.relative_to(include_root)}:{line_number}: {include}")
+
+    if failures:
+        raise RuntimeError("Missing local NCCL EP JIT headers:\n" + "\n".join(failures))
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -208,7 +208,7 @@ def _discover_nccl_home() -> str:
 
 
 def build_nccl_ep_submodule() -> str:
-    """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME."""
+    """Build NCCL EP libraries from 3rdparty/nccl-extensions and return NCCL_HOME."""
     nccl_root = current_file_path / "3rdparty" / "nccl-extensions"
     if not (nccl_root / "nccl_ep" / "Makefile").exists():
         raise RuntimeError(
@@ -217,7 +217,7 @@ def build_nccl_ep_submodule() -> str:
         )
 
     build_dir = nccl_root / "build"
-    nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a"
+    nccl_ep_shared_lib = build_dir / "lib" / "libnccl_ep.so"
     gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode"
 
     # Caller gates on arch >= 90 or "native"; expand "native" to the host's
@@ -261,18 +261,18 @@ def build_nccl_ep_submodule() -> str:
     env["NCCL_EP_BUILDDIR"] = str(build_dir)
 
     prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None
-    if not nccl_ep_lib.exists() or prev_gencode != gencode:
-        if nccl_ep_lib.exists() and prev_gencode != gencode:
+    if not nccl_ep_shared_lib.exists() or prev_gencode != gencode:
+        if nccl_ep_shared_lib.exists() and prev_gencode != gencode:
             print(
                 f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); "
-                "rebuilding libnccl_ep.a"
+                "rebuilding NCCL EP libraries"
             )
             subprocess.check_call(
                 ["make", "-C", "nccl_ep", "clean"],
                 cwd=str(nccl_root),
                 env=env,
             )
-        print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')")
+        print(f"[NCCL EP] Building static and shared libraries (gencode='{gencode}')")
         make_jobs = f"-j{nproc}" if nproc else "-j"
         subprocess.check_call(
             ["make", make_jobs, "-C", "nccl_ep", "lib"],

tests/cpp_distributed/CMakeLists.txt

@@ -101,7 +101,7 @@ gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600)
 # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h).
 # The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*);
 # all ncclEp* calls live behind TE's public <transformer_engine/ep.h>, which is
-# statically linked into libtransformer_engine.so.
+# backed by the runtime-loaded libnccl_ep.so library.
 message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}")
 set(EP_TEST_COMMON_INCLUDES
     ${NCCL_INCLUDE_DIR}
@@ -129,4 +129,4 @@ target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS})
 
 # Do NOT use gtest_discover_tests - these binaries require multi-process
 # launch via run_test_ep.sh, not direct single-process execution.
-message(STATUS "EP distributed tests enabled (NCCL EP statically linked into libtransformer_engine.so)")
+message(STATUS "EP distributed tests enabled (NCCL EP loaded from libnccl_ep.so at runtime)")

tests/cpp_distributed/run_test_ep.sh

@@ -53,6 +53,12 @@ if (( NUM_GPUS < 2 )); then
     exit 0
 fi
 
+# Force this test run to compile at least one NCCL EP JIT kernel instead of
+# succeeding from a cache populated by an earlier process.
+NCCL_EP_JIT_CACHE_DIR=$(mktemp -d "${TMPDIR:-/tmp}/te-nccl-ep-jit.XXXXXX")
+export NCCL_EP_JIT_CACHE_DIR
+trap 'rm -rf "${NCCL_EP_JIT_CACHE_DIR}"' EXIT
+
 GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}"
 
 echo "=== EP Tests ==="
@@ -67,3 +73,8 @@ if [[ -n "${GTEST_XML_PREFIX:-}" ]]; then
 else
     "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS}
 fi
+
+if ! compgen -G "${NCCL_EP_JIT_CACHE_DIR}/*/*.cubin" > /dev/null; then
+    echo "ERROR: NCCL EP tests did not produce a JIT-compiled cubin."
+    exit 1
+fi

tests/pytorch/test_nccl_ep_discovery.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import transformer_engine
+
+
+def _hide_packaged_library(monkeypatch):
+    def _missing_library(_):
+        raise FileNotFoundError
+
+    monkeypatch.setattr(
+        transformer_engine.common,
+        "_get_shared_object_file",
+        _missing_library,
+    )
+
+
+def test_nccl_ep_library_found_from_home(monkeypatch, tmp_path):
+    home = tmp_path / "nccl_ep"
+    library_dir = home / "lib"
+    library_dir.mkdir(parents=True)
+    (library_dir / "libnccl_ep.so.0.1").touch()
+
+    monkeypatch.setenv("NCCL_EP_HOME", str(home))
+    _hide_packaged_library(monkeypatch)
+    monkeypatch.setattr(transformer_engine, "find_library", lambda _: None)
+
+    assert transformer_engine._nccl_ep_library_installed()
+
+
+def test_nccl_ep_library_found_in_package(monkeypatch, tmp_path):
+    library = tmp_path / "libnccl_ep.so"
+    library.touch()
+
+    monkeypatch.delenv("NCCL_EP_HOME", raising=False)
+    monkeypatch.setattr(
+        transformer_engine.common,
+        "_get_shared_object_file",
+        lambda _: library,
+    )
+    monkeypatch.setattr(transformer_engine, "find_library", lambda _: None)
+
+    assert transformer_engine._nccl_ep_library_installed()
+
+
+def test_nccl_ep_library_found_by_dynamic_loader(monkeypatch):
+    monkeypatch.delenv("NCCL_EP_HOME", raising=False)
+    _hide_packaged_library(monkeypatch)
+    monkeypatch.setattr(transformer_engine, "find_library", lambda _: "libnccl_ep.so.0")
+
+    assert transformer_engine._nccl_ep_library_installed()
+
+
+def test_nccl_ep_library_not_found(monkeypatch):
+    monkeypatch.delenv("NCCL_EP_HOME", raising=False)
+    _hide_packaged_library(monkeypatch)
+    monkeypatch.setattr(transformer_engine, "find_library", lambda _: None)
+
+    assert not transformer_engine._nccl_ep_library_installed()

transformer_engine/__init__.py

@@ -9,11 +9,13 @@
 import ctypes
 import functools
 import os
+from ctypes.util import find_library
 from importlib import metadata
+from pathlib import Path
 from typing import Optional, Tuple
 import transformer_engine.common
 
-# Minimum NCCL version for the statically-linked NCCL EP backend.
+# Minimum NCCL version for the runtime-loaded NCCL EP backend.
 _NCCL_EP_MIN_VERSION = (2, 30, 4)
 
 
@@ -32,14 +34,39 @@ def _nccl_runtime_version() -> Optional[Tuple[int, int, int]]:
     return (v // 10000, (v // 100) % 100, v % 100)
 
 
+def _nccl_ep_library_installed() -> bool:
+    if nccl_ep_home := os.getenv("NCCL_EP_HOME"):
+        home = Path(nccl_ep_home)
+        if any(
+            library.is_file()
+            for lib_dir in ("lib", "lib64")
+            for library in (home / lib_dir).glob("libnccl_ep.so*")
+        ):
+            return True
+
+    try:
+        transformer_engine.common._get_shared_object_file("nccl_ep")
+    except FileNotFoundError:
+        return find_library("nccl_ep") is not None
+    else:
+        return True
+
+
 def is_nccl_ep_available() -> bool:
-    """Return True if the runtime libnccl.so meets the NCCL EP minimum."""
+    """Return True if the NCCL EP library and a compatible NCCL runtime are available."""
+    if not _nccl_ep_library_installed():
+        return False
     cur = _nccl_runtime_version()
     return cur is not None and cur >= _NCCL_EP_MIN_VERSION
 
 
 def require_nccl_ep() -> None:
     """Raise RuntimeError if NCCL EP cannot run on the current libnccl."""
+    if not _nccl_ep_library_installed():
+        raise RuntimeError(
+            "NCCL EP library libnccl_ep.so is not installed. Build Transformer Engine "
+            "with NVTE_WITH_NCCL_EP=1."
+        )
     mn = ".".join(str(x) for x in _NCCL_EP_MIN_VERSION)
     cur = _nccl_runtime_version()
     if cur is None:

transformer_engine/common/CMakeLists.txt

@@ -455,7 +455,7 @@ endif()
 # Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to
 # skip NCCL EP entirely - useful on older images whose system NCCL is below
 # the 2.30.4 EP minimum.
-option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON)
+option(NVTE_WITH_NCCL_EP "Build and package the runtime-loaded NCCL EP library" ON)
 if(NVTE_WITH_NCCL_EP)
 # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize.
 # -- NCCL EP headers --------------------------------------------------------
@@ -471,17 +471,17 @@ if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h")
 endif()
 message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}")
 
-# -- libnccl_ep.a -----------------------------------------------------------
-# Statically linked into libtransformer_engine.so. EPBackend::initialize checks
-# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep
-# imports stay unresolved (and harmless) under default ELF lazy binding when
-# the gate trips. LD_BIND_NOW environments lose this property.
+# -- libnccl_ep.so ----------------------------------------------------------
+# Loaded on demand by EPBackend only after it verifies NCCL >= 2.30.4. Keeping
+# NCCL EP out of libtransformer_engine.so allows non-EP users to run with older
+# NCCL versions even when the core library uses eager symbol binding.
 set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib")
-find_file(NCCL_EP_LIB
-    NAMES libnccl_ep.a
+find_file(NCCL_EP_SHARED_LINK
+    NAMES libnccl_ep.so
     HINTS ${NCCL_EP_LIB_DIR}
     NO_DEFAULT_PATH
     REQUIRED)
+file(REAL_PATH "${NCCL_EP_SHARED_LINK}" NCCL_EP_SHARED_LIB)
 
 # -- NCCL core library -------------------------------------------------------
 if(NOT NCCL_LIB)
@@ -494,20 +494,17 @@ endif()
 target_include_directories(transformer_engine PRIVATE
     ${NCCL_EP_INCLUDE_DIR})
 
-# libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's
-# DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow
-# the static archive on the link line so --as-needed records libcuda.so.1.
+# libtransformer_engine.so uses only the baseline NCCL API directly. NCCL EP's
+# newer symbols remain isolated in libnccl_ep.so until EP is initialized.
 target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB})
-target_link_libraries(transformer_engine PRIVATE
-    -Wl,--whole-archive ${NCCL_EP_LIB} -Wl,--no-whole-archive
-    CUDA::cuda_driver)
 
 target_sources(transformer_engine PRIVATE
     ep/ep_backend.cpp
-    ep/ep_api.cpp)
+    ep/ep_api.cpp
+    ep/nccl_ep_provider.cpp)
 target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP)
 
-message(STATUS "NCCL EP enabled (static link): ${NCCL_EP_LIB}")
+message(STATUS "NCCL EP enabled (runtime-loaded library): ${NCCL_EP_SHARED_LIB}")
 message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}")
 else()
   # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs.
@@ -677,3 +674,16 @@ message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}")
 
 # Install library
 install(TARGETS transformer_engine DESTINATION .)
+if(NVTE_WITH_NCCL_EP)
+  # Install a real file rather than the build-tree symlink so wheel packaging
+  # does not need to preserve the versioned symlink chain.
+  install(FILES "${NCCL_EP_SHARED_LIB}"
+          DESTINATION .
+          RENAME libnccl_ep.so)
+  # NCCL EP JIT compilation requires both the public header root and staged
+  # device headers under NCCL_EP_HOME/include/nccl_ep.
+  install(FILES "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h"
+          DESTINATION nccl_ep/include)
+  install(DIRECTORY "${NCCL_EP_INCLUDE_DIR}/nccl_ep"
+          DESTINATION nccl_ep/include)
+endif()

transformer_engine/common/__init__.py

@@ -107,9 +107,11 @@ def _get_shared_object_file(library: str) -> Path:
     """
 
     # Check provided input and determine the correct prefix for .so.
-    assert library in ("core", "torch", "jax"), f"Unsupported TE library {library}."
+    assert library in ("core", "torch", "jax", "nccl_ep"), f"Unsupported TE library {library}."
     if library == "core":
         so_prefix = "libtransformer_engine"
+    elif library == "nccl_ep":
+        so_prefix = "libnccl_ep"
     else:
         so_prefix = f"transformer_engine_{library}"