From c03502001bcaebb0d5a894e9036570128505b13e Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 27 Aug 2025 13:44:43 -0700 Subject: [PATCH 1/3] Fix CUDA version in setup.py Signed-off-by: Vladimir Cherepanov --- setup.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 52adaf9238..ed1f5b8a9d 100644 --- a/setup.py +++ b/setup.py @@ -17,6 +17,7 @@ from build_tools.te_version import te_version from build_tools.utils import ( cuda_archs, + cuda_version, get_frameworks, remove_dups, ) @@ -70,11 +71,11 @@ def setup_common_extension() -> CMakeExtension: if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))): cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON") cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution( - "nvidia-cublasmp-cu12" - ).locate_file("nvidia/cublasmp/cu12") + f"nvidia-cublasmp-cu{cuda_version()[0]}" + ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}") cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}") nvshmem_dir = os.getenv("NVSHMEM_HOME") or metadata.distribution( - "nvidia-nvshmem-cu12" + f"nvidia-nvshmem-cu{cuda_version()[0]}" ).locate_file("nvidia/nvshmem") cmake_flags.append(f"-DNVSHMEM_DIR={nvshmem_dir}") print("CMAKE_FLAGS:", cmake_flags[-2:]) From e5fcc2cb0ae1f861b10aaef262dbeb6102c627cc Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 28 Aug 2025 18:09:02 +0000 Subject: [PATCH 2/3] Re-enable building comm-gemm tests Signed-off-by: Vladimir Cherepanov --- tests/cpp/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index c2c9d0d915..412c5d34d9 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -43,5 +43,6 @@ include_directories(${CMAKE_SOURCE_DIR}) find_package(CUDAToolkit REQUIRED) include(${CMAKE_SOURCE_DIR}/../../3rdparty/cudnn-frontend/cmake/cuDNN.cmake) +add_subdirectory(comm_gemm) add_subdirectory(operator) add_subdirectory(util) From b1a9aa9663c768963f3cc39d81e31213a4675a62 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 4 Sep 2025 14:16:39 -0700 Subject: [PATCH 3/3] WAR for nvidia-nvshmem package Signed-off-by: Vladimir Cherepanov --- transformer_engine/common/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 7feb5fda5f..dd1ec480b2 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -218,6 +218,11 @@ def _nvidia_cudart_include_dir() -> str: except ModuleNotFoundError: return "" + # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia" + # above doesn't through. However, they don't set "__file__" attribute. + if nvidia.__file__ is None: + return "" + include_dir = Path(nvidia.__file__).parent / "cuda_runtime" return str(include_dir) if include_dir.exists() else ""