From 44740c832f79dc06db55fceb95ad10ab5438a63c Mon Sep 17 00:00:00 2001
From: Kevin Tong <kevin@augmentcode.com>
Date: Tue, 14 Oct 2025 21:57:52 +0000
Subject: [PATCH] NVTX Patch

---
 transformer_engine/CMakeLists.txt        | 2 +-
 transformer_engine/common/CMakeLists.txt | 1 -
 transformer_engine/common/nvtx.h         | 2 +-
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/CMakeLists.txt b/transformer_engine/CMakeLists.txt
index 91e98c835d..9b1d012e9d 100644
--- a/transformer_engine/CMakeLists.txt
+++ b/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)
diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
index 7426c3112f..fe6453f774 100644
--- a/transformer_engine/common/CMakeLists.txt
+++ b/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})
diff --git a/transformer_engine/common/nvtx.h b/transformer_engine/common/nvtx.h
index db51fc1a08..009561d68a 100644
--- a/transformer_engine/common/nvtx.h
+++ b/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 {