Fix error when building with WITH_NVSHMEM=ON (#73368)

XieYunshen · web-flow · commit 0c8a18fb64f5 · 2025-06-19T16:24:01.000+08:00
* Fix error when building with WITH_NVSHMEM=ON

* fix cannot import deep_ep error

* auto-format code via pre-commit
diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake
@@ -17,13 +17,9 @@ endif()
 
 find_path(
   CUDNN_INCLUDE_DIR cudnn.h
-  PATHS ${CUDNN_ROOT}
-        ${CUDNN_ROOT}/include
-        ${CUDNN_ROOT}/include/${TARGET_ARCH}-linux-gnu
-        $ENV{CUDNN_ROOT}
-        $ENV{CUDNN_ROOT}/include
-        ${CUDA_TOOLKIT_INCLUDE}
-        /usr/local/lib/python${PY_VERSION}/dist-packages/nvidia/cudnn/include/
+  PATHS ${CUDNN_ROOT} ${CUDNN_ROOT}/include
+        ${CUDNN_ROOT}/include/${TARGET_ARCH}-linux-gnu $ENV{CUDNN_ROOT}
+        $ENV{CUDNN_ROOT}/include ${CUDA_TOOLKIT_INCLUDE}
   NO_DEFAULT_PATH)
 
 get_filename_component(__libpath_hist ${CUDA_CUDART_LIBRARY} PATH)
@@ -66,7 +62,6 @@ find_library(
   CUDNN_LIBRARY
   NAMES ${CUDNN_LIB_NAME} # libcudnn_static.a
   PATHS ${CUDNN_CHECK_LIBRARY_DIRS} ${CUDNN_INCLUDE_DIR} ${__libpath_hist}
-        /usr/local/lib/python${PY_VERSION}/dist-packages/nvidia/cudnn/lib/
   NO_DEFAULT_PATH
   DOC "Path to cuDNN library.")
 
diff --git a/cmake/external/nvshmem.cmake b/cmake/external/nvshmem.cmake
@@ -55,9 +55,9 @@ endif()
 
 set(NVSHMEM_PATCH_PATH ${PADDLE_SOURCE_DIR}/patches/nvshmem/nvshmem.patch)
 set(NVSHMEM_PATCH_COMMAND
-    git init && git config user.name "PaddlePaddle" && git config user.email
-    "paddle@baidu.com" && git config --global --add safe.directory
-    ${NVSHMEM_SOURCE_DIR} && git add . && git commit -m "init" && git apply
+    git init && git config --global --add safe.directory ${NVSHMEM_SOURCE_DIR}
+    && git config user.name "PaddlePaddle" && git config user.email
+    "paddle@baidu.com" && git add . && git commit -m "init" && git apply
     ${NVSHMEM_PATCH_PATH})
 
 set(NVSHMEM_LIB ${NVSHMEM_INSTALL_DIR}/lib/libnvshmem.a)
diff --git a/setup.py b/setup.py
@@ -2432,8 +2432,8 @@ def get_setup_parameters():
 
     if (
         env_dict.get("WITH_GPU") == 'ON'
-        and env_dict.get("CUDA_ARCH_BIN")
-        and env_dict.get("CUDA_ARCH_BIN").find("90") != -1
+        and env_dict.get("COMPILED_CUDA_ARCHS")
+        and env_dict.get("COMPILED_CUDA_ARCHS").find("90") != -1
     ):
         packages.extend(['paddle.distributed.communication.deep_ep'])
     if (
