diff --git a/.gitmodules b/.gitmodules index 495d8e3fe7..07647e915e 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,9 +1,6 @@ [submodule "3rdparty/googletest"] path = 3rdparty/googletest url = https://github.com/google/googletest.git -[submodule "3rdparty/cudnn-frontend"] - path = 3rdparty/cudnn-frontend - url = https://github.com/NVIDIA/cudnn-frontend.git [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend deleted file mode 160000 index e46d708245..0000000000 --- a/3rdparty/cudnn-frontend +++ /dev/null @@ -1 +0,0 @@ -Subproject commit e46d7082450264ce05cf898f8740011c4896f817 diff --git a/MANIFEST.in b/MANIFEST.in index c2309a0370..f97ab40cb8 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,3 @@ recursive-include transformer_engine/common/include *.* +recursive-include transformer_engine/common/cmake *.cmake recursive-include build_tools *.py *.txt diff --git a/build_tools/jax.py b/build_tools/jax.py index 654f8bfc17..d61c26c128 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -13,6 +13,7 @@ from .utils import ( get_cuda_include_dirs, all_files_in_dir, + cudnn_frontend_include_path, debug_build_enabled, setup_mpi_flags, nccl_ep_enabled, @@ -22,7 +23,7 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax", "flax>=0.7.1"] + return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] def test_requirements() -> List[str]: @@ -89,20 +90,7 @@ def setup_jax_extension( # Header files include_dirs = get_cuda_include_dirs() - cudnn_frontend_include_dir = None - for base_path in (Path(common_header_files), *Path(common_header_files).parents): - candidate = base_path / "3rdparty" / "cudnn-frontend" / "include" - if candidate.exists(): - cudnn_frontend_include_dir = candidate - break - if cudnn_frontend_include_dir is None: - for base_path in Path(__file__).resolve().parents: - candidate = base_path / "3rdparty" / "cudnn-frontend" / "include" - if candidate.exists(): - cudnn_frontend_include_dir = candidate - break - if cudnn_frontend_include_dir is not None: - include_dirs.append(cudnn_frontend_include_dir) + include_dirs.append(cudnn_frontend_include_path()) include_dirs.extend( [ common_header_files, diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 5ed4eae9d5..2bb238c522 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -22,7 +22,16 @@ def install_requirements() -> List[str]: """Install dependencies for TE/PyTorch extensions.""" - return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic", "nvdlfw-inspect"] + return [ + "torch>=2.1", + "einops", + "onnxscript", + "onnx", + "packaging", + "pydantic", + "nvdlfw-inspect", + "nvidia-cudnn-frontend>=1.25.0", + ] def test_requirements() -> List[str]: diff --git a/build_tools/utils.py b/build_tools/utils.py index a3bb8c23f2..6a8168e7d0 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -14,7 +14,7 @@ import sys import platform from pathlib import Path -from importlib.metadata import version as get_version +from importlib.metadata import PackageNotFoundError, distribution, version as get_version from subprocess import CalledProcessError from typing import List, Optional, Tuple, Union @@ -251,6 +251,26 @@ def get_cuda_include_dirs() -> Tuple[str, str]: ] +@functools.lru_cache(maxsize=None) +def cudnn_frontend_include_path() -> Path: + """Return the C++ include directory from nvidia-cudnn-frontend.""" + package = "nvidia-cudnn-frontend" + try: + include_dir = Path(distribution(package).locate_file("include")).resolve() + except PackageNotFoundError as e: + raise RuntimeError( + f"{package} is required to build Transformer Engine. " + f"Install it with `pip install {package}`." + ) from e + + header = include_dir / "cudnn_frontend.h" + if not header.is_file(): + raise RuntimeError( + f"The {package} installation does not contain the expected header {header}." + ) + return include_dir + + @functools.lru_cache(maxsize=None) def cuda_archs() -> str: archs = os.getenv("NVTE_CUDA_ARCHS") diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index ff422e9d2c..8fe2b629be 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -23,7 +23,7 @@ git checkout $TARGET_BRANCH git submodule update --init --recursive # Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' if $BUILD_METAPACKAGE ; then cd /TransformerEngine diff --git a/pyproject.toml b/pyproject.toml index 4a8fded172..2c9f224c14 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1"] +requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" diff --git a/setup.py b/setup.py index 1eead737d5..150d92969c 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ from build_tools.utils import ( cuda_archs, cuda_version, + cudnn_frontend_include_path, get_frameworks, remove_dups, min_python_version_str, @@ -56,7 +57,10 @@ def run(self): def setup_common_extension() -> CMakeExtension: """Setup CMake extension for common library""" - cmake_flags = ["-DCMAKE_CUDA_ARCHITECTURES={}".format(archs)] + cmake_flags = [ + "-DCMAKE_CUDA_ARCHITECTURES={}".format(archs), + f"-DCUDNN_FRONTEND_INCLUDE_DIR={cudnn_frontend_include_path()}", + ] if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): assert ( os.getenv("MPI_HOME") is not None diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 6f4f163f08..2d5469971b 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -41,7 +41,7 @@ include_directories(../../transformer_engine) include_directories(${CMAKE_SOURCE_DIR}) find_package(CUDAToolkit REQUIRED) -include(${CMAKE_SOURCE_DIR}/../../3rdparty/cudnn-frontend/cmake/cuDNN.cmake) +include(${CMAKE_SOURCE_DIR}/../../transformer_engine/common/cmake/cuDNN.cmake) add_subdirectory(operator) add_subdirectory(util) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index edb8c5e109..be64fcb2be 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -88,15 +88,12 @@ endif() set(NVTE_STANDARD_ARCHS ${CMAKE_CUDA_ARCHITECTURES}) # cuDNN frontend API -set(CUDNN_FRONTEND_INCLUDE_DIR - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cudnn-frontend/include") -if(NOT EXISTS "${CUDNN_FRONTEND_INCLUDE_DIR}") +if(NOT EXISTS "${CUDNN_FRONTEND_INCLUDE_DIR}/cudnn_frontend.h") message(FATAL_ERROR - "Could not find cuDNN frontend API at ${CUDNN_FRONTEND_INCLUDE_DIR}. " - "Try running 'git submodule update --init --recursive' " - "within the Transformer Engine source.") + "Could not find cudnn_frontend.h in ${CUDNN_FRONTEND_INCLUDE_DIR}. " + "Install nvidia-cudnn-frontend or set CUDNN_FRONTEND_INCLUDE_DIR.") endif() -include(${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cudnn-frontend/cmake/cuDNN.cmake) +include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/cuDNN.cmake) set(CUTLASS_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cutlass/include") diff --git a/transformer_engine/common/cmake/cuDNN.cmake b/transformer_engine/common/cmake/cuDNN.cmake new file mode 100644 index 0000000000..80f0855b30 --- /dev/null +++ b/transformer_engine/common/cmake/cuDNN.cmake @@ -0,0 +1,163 @@ +# Check if CUDNN:: targets already exist (integrated build) +if(TARGET CUDNN::cudnn AND TARGET CUDNN::cudnn_all) + message(STATUS "cuDNN: Using existing CMake targets") + set(CUDNN_FOUND ON CACHE INTERNAL "cuDNN Library Found") + return() +endif() + +message(STATUS "cuDNN: Searching for pre-built libraries") + +add_library(CUDNN::cudnn_all INTERFACE IMPORTED) + +find_path( + CUDNN_INCLUDE_DIR cudnn.h + HINTS $ENV{CUDNN_INCLUDE_PATH} ${CUDNN_INCLUDE_PATH} $ENV{CUDNN_PATH} ${CUDNN_PATH} ${Python_SITEARCH}/nvidia/cudnn ${CUDAToolkit_INCLUDE_DIRS} + PATH_SUFFIXES include +) + +if(CUDNN_INCLUDE_DIR) + if(EXISTS ${CUDNN_INCLUDE_DIR}/cudnn_version.h) + file(READ ${CUDNN_INCLUDE_DIR}/cudnn_version.h CUDNN_HEADER_CONTENTS) + else() + file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_HEADER_CONTENTS) + endif() + string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)" + CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}") + string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)" + CUDNN_VERSION_MINOR "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}") + string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)" + CUDNN_VERSION_PATCH "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1" + CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}") + if(NOT CUDNN_VERSION_MAJOR) + set(CUDNN_VERSION "?") + else() + set(CUDNN_VERSION + "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}") + endif() + set(CUDNN_MAJOR_VERSION ${CUDNN_VERSION_MAJOR}) +endif() + +function(find_cudnn_library NAME) + if(NOT "${ARGV1}" STREQUAL "OPTIONAL") + set(_cudnn_required "REQUIRED") + else() + set(_cudnn_required "") + endif() + + if(CUDNN_STATIC) + set(library_names "${NAME}_static" "lib${NAME}_static_v${CUDNN_MAJOR_VERSION}.a") + else() + set(library_names ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}") + endif() + find_library( + ${NAME}_LIBRARY + NAMES ${library_names} + NAMES_PER_DIR + HINTS $ENV{CUDNN_LIBRARY_PATH} ${CUDNN_LIBRARY_PATH} $ENV{CUDNN_PATH} ${CUDNN_PATH} ${Python_SITEARCH}/nvidia/cudnn ${CUDAToolkit_LIBRARY_DIR} + PATH_SUFFIXES lib64 lib/x64 lib + ${_cudnn_required} + ) + + if(${NAME}_LIBRARY) + add_library(CUDNN::${NAME} UNKNOWN IMPORTED) + set_target_properties( + CUDNN::${NAME} PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR} + IMPORTED_LOCATION ${${NAME}_LIBRARY} + ) + message(STATUS "${NAME} found at ${${NAME}_LIBRARY}.") + else() + message(STATUS "${NAME} not found.") + endif() +endfunction() + +if(NOT CUDNN_STATIC) + find_cudnn_library(cudnn) + set(CUDNN_LIBRARY_VAR cudnn_LIBRARY) +endif() + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args( + LIBRARY REQUIRED_VARS + CUDNN_INCLUDE_DIR ${CUDNN_LIBRARY_VAR} +) + +if(CUDNN_INCLUDE_DIR AND (CUDNN_STATIC OR cudnn_LIBRARY)) + message(STATUS "cuDNN: ${cudnn_LIBRARY}") + message(STATUS "cuDNN: ${CUDNN_INCLUDE_DIR}") + set(CUDNN_FOUND ON CACHE INTERNAL "cuDNN Library Found") +else() + set(CUDNN_FOUND OFF CACHE INTERNAL "cuDNN Library Not Found") +endif() + +target_include_directories( + CUDNN::cudnn_all + INTERFACE + $ + $ +) + +if(CUDNN_STATIC) + add_library(CUDNN::cudnn INTERFACE IMPORTED) + target_link_libraries( + CUDNN::cudnn + INTERFACE + CUDNN::cudnn_all + ) + find_package(ZLIB REQUIRED) +else() + target_link_libraries( + CUDNN::cudnn_all + INTERFACE + CUDNN::cudnn + ) +endif() + +if(CUDNN_MAJOR_VERSION EQUAL 8) + find_cudnn_library(cudnn_adv_infer) + find_cudnn_library(cudnn_adv_train) + find_cudnn_library(cudnn_cnn_infer) + find_cudnn_library(cudnn_cnn_train) + find_cudnn_library(cudnn_ops_infer) + find_cudnn_library(cudnn_ops_train) + + target_link_libraries( + CUDNN::cudnn_all + INTERFACE + CUDNN::cudnn_adv_train + CUDNN::cudnn_ops_train + CUDNN::cudnn_cnn_train + CUDNN::cudnn_adv_infer + CUDNN::cudnn_cnn_infer + CUDNN::cudnn_ops_infer + ) +elseif(CUDNN_MAJOR_VERSION EQUAL 9) + find_cudnn_library(cudnn_graph) + find_cudnn_library(cudnn_engines_runtime_compiled) + find_cudnn_library(cudnn_ops OPTIONAL) + find_cudnn_library(cudnn_cnn OPTIONAL) + find_cudnn_library(cudnn_adv OPTIONAL) + find_cudnn_library(cudnn_engines_precompiled OPTIONAL) + find_cudnn_library(cudnn_heuristic OPTIONAL) + find_cudnn_library(cudnn_ext OPTIONAL) + + target_link_libraries( + CUDNN::cudnn_all + INTERFACE + $<$:-Wl,--whole-archive> + CUDNN::cudnn_graph + CUDNN::cudnn_engines_runtime_compiled + CUDNN::cudnn_ops + CUDNN::cudnn_cnn + CUDNN::cudnn_adv + $<$>:CUDNN::cudnn_engines_precompiled> + CUDNN::cudnn_heuristic + $<$:-Wl,--no-whole-archive> + $<$:CUDA::cublasLt_static $,CUDA::nvrtc_static,CUDA::nvrtc> ZLIB::ZLIB> + ) +endif() diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index 71f0176ed7..ff9f5fb146 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -2,12 +2,11 @@ # # See LICENSE for license information. """cuDNN frontend score_mod fused attention helpers.""" + import hashlib import importlib import inspect import os -from pathlib import Path -import sys from dataclasses import dataclass from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple @@ -18,7 +17,6 @@ import transformer_engine_jax - __all__ = [ "FusedAttnScoreModHelper", "make_fused_attn_score_mod_config", @@ -27,10 +25,6 @@ "fused_attn_score_mod_bwd", ] -_CUDNN_FRONTEND_PYTHON_PATH = ( - Path(__file__).resolve().parents[3] / "3rdparty" / "cudnn-frontend" / "python" -) - def _is_non_deterministic_allowed(): """Check if non-deterministic kernels are allowed.""" @@ -600,13 +594,6 @@ def _shape_dtype(value) -> jax.ShapeDtypeStruct: def _import_cudnn_for_score_mod(): - cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH) - cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn" - if ( - any(cudnn_frontend_package.glob("_compiled_module*")) - and cudnn_frontend_path not in sys.path - ): - sys.path.insert(0, cudnn_frontend_path) try: cudnn = importlib.import_module("cudnn") except ImportError as exc: diff --git a/transformer_engine/jax/pyproject.toml b/transformer_engine/jax/pyproject.toml index d3162ae96d..d5c15b4a8f 100755 --- a/transformer_engine/jax/pyproject.toml +++ b/transformer_engine/jax/pyproject.toml @@ -3,8 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "pybind11[global]", "pip", "jax[cuda12]", "flax>=0.7.1"] +requires = ["setuptools>=61.0", "pybind11[global]", "pip", "jax[cuda12]", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" - diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index da8205923f..b9593b42d9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -6,10 +6,7 @@ from dataclasses import dataclass import importlib -import importlib.util import inspect -from pathlib import Path -import sys from typing import Any, Callable, Dict, Optional, Tuple import torch @@ -17,27 +14,17 @@ _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} _SCORE_MOD_UNCACHEABLE = object() -_CUDNN_FRONTEND_PYTHON_PATH = ( - Path(__file__).resolve().parents[4] / "3rdparty" / "cudnn-frontend" / "python" -) def _import_cudnn_frontend(): - """Import the vendored cuDNN frontend if built, otherwise use the installed package.""" - cudnn_frontend_path = str(_CUDNN_FRONTEND_PYTHON_PATH) - cudnn_frontend_package = _CUDNN_FRONTEND_PYTHON_PATH / "cudnn" - if any(cudnn_frontend_package.glob("_compiled_module*")): - if cudnn_frontend_path not in sys.path: - sys.path.insert(0, cudnn_frontend_path) - return importlib.import_module("cudnn") - - if importlib.util.find_spec("cudnn") is not None: + """Import the cuDNN frontend Python package.""" + try: return importlib.import_module("cudnn") - - raise ImportError( - "cuDNN frontend Python package not found. " - "Install it with: pip install nvidia-cudnn-frontend" - ) + except ImportError as exc: + raise ImportError( + "cuDNN frontend Python package not found. " + "Install it with: pip install nvidia-cudnn-frontend" + ) from exc def _bhsd_dim_stride(