Remove cuBLASMp dir hint

Signed-off-by: Vladimir Cherepanov <[email protected]>

Author: vcherepanov-nv

setup.py

@@ -70,10 +70,6 @@ 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(
-            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(
             f"nvidia-nvshmem-cu{cuda_version()[0]}"
         ).locate_file("nvidia/nvshmem")

tests/cpp/comm_gemm/CMakeLists.txt

@@ -12,7 +12,7 @@ find_library(NCCL_LIB
              NAMES nccl libnccl
              PATH_SUFFIXES lib
              REQUIRED)
-target_include_directories(test_comm_gemm PRIVATE ${MPI_CXX_INCLUDE_PATH} $ENV{CUBLASMP_HOME}/include)
+target_include_directories(test_comm_gemm PRIVATE ${MPI_CXX_INCLUDE_PATH})
 target_link_libraries(test_comm_gemm PUBLIC CUDA::cuda_driver CUDA::cudart GTest::gtest ${TE_LIB} CUDA::nvrtc CUDNN::cudnn MPI::MPI_CXX ${NCCL_LIB} OpenMP::OpenMP_CXX)
 
 include(GoogleTest)

transformer_engine/common/CMakeLists.txt

@@ -152,10 +152,9 @@ endif()
 option(NVTE_WITH_CUBLASMP "Use cuBLASMp for tensor parallel GEMMs" OFF)
 if (NVTE_WITH_CUBLASMP)
     target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP)
-    target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include ${NVSHMEM_DIR}/include)
+    target_include_directories(transformer_engine PRIVATE ${NVSHMEM_DIR}/include)
     find_library(CUBLASMP_LIB
                  NAMES cublasmp libcublasmp
-                 PATHS ${CUBLASMP_DIR}
                  PATH_SUFFIXES lib
                  REQUIRED)
     find_library(NVSHMEM_HOST_LIB
@@ -164,8 +163,6 @@ if (NVTE_WITH_CUBLASMP)
                  PATH_SUFFIXES lib
                  REQUIRED)
   target_link_libraries(transformer_engine PUBLIC ${CUBLASMP_LIB} ${NVSHMEM_HOST_LIB})
-  message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}")
-  message(STATUS "Using nvshmem at: ${NVSHMEM_DIR}")
 endif()
 
 # Hack to enable dynamic loading in cuDNN frontend