NVTX Patch

Author: kevin-tong-augment

transformer_engine/CMakeLists.txt

@@ -20,7 +20,7 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug")
 endif()
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}/cmake/")
-find_package(CUDAToolkit REQUIRED cublas nvToolsExt)
+find_package(CUDAToolkit REQUIRED cublas)
 find_package(CUDNN REQUIRED cudnn)
 # NOTE[augment] -- commenting out this line since the Python components are unused and they cause build issues
 # find_package(Python COMPONENTS Interpreter Development REQUIRED)

transformer_engine/common/CMakeLists.txt

@@ -55,7 +55,6 @@ target_link_libraries(transformer_engine PUBLIC
                       CUDA::cuda_driver
                       CUDA::cudart
                       CUDA::nvrtc
-                      CUDA::nvToolsExt
                       CUDNN::cudnn)
 target_include_directories(transformer_engine PRIVATE
                            ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})

transformer_engine/common/nvtx.h

@@ -8,7 +8,7 @@
 #define TRANSFORMER_ENGINE_COMMON_NVTX_H_
 
 #include <string>
-#include <nvToolsExt.h>
+#include <nvtx3/nvToolsExt.h>
 
 namespace transformer_engine::nvtx {