Make it an option to compile activation functions with fast math (#1410)

* Add a compile option to compile activation kernels with fast math Signed-off-by: Guyue Huang <[email protected]> * Fix Signed-off-by: Guyue Huang <[email protected]> * Apply suggestions from code review Co-authored-by: Kirthi Shankar Sivamani <[email protected]> Signed-off-by: guyueh1 <[email protected]> --------- Signed-off-by: Guyue Huang <[email protected]> Signed-off-by: guyueh1 <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
NVIDIA · Jan 15, 2025 · 3d63cbb · 3d63cbb
1 parent 2402406
commit 3d63cbb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -64,6 +64,9 @@ def setup_common_extension() -> CMakeExtension:
         ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1"
         cmake_flags.append("-DNVTE_UB_WITH_MPI=ON")
 
+    if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))):
+        cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON")
+
     # Project directory root
     root_path = Path(__file__).resolve().parent
 

diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
@@ -147,6 +147,14 @@ set_source_files_properties(fused_softmax/scaled_masked_softmax.cu
                             fused_softmax/scaled_aligned_causal_masked_softmax.cu
                             PROPERTIES
                             COMPILE_OPTIONS "--use_fast_math")
+option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF)
+if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH)
+  set_source_files_properties(activation/gelu.cu
+                              activation/relu.cu
+                              activation/swiglu.cu
+                              PROPERTIES
+                              COMPILE_OPTIONS "--use_fast_math")
+endif()
 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -O3")