From af67bf00a894fa82abedd70539e5337e7d5584ad Mon Sep 17 00:00:00 2001
From: William Tambellini <wtambellini@sdl.com>
Date: Sun, 13 Apr 2025 19:30:23 -0700
Subject: [PATCH] feat: register ENABLE_MULTI_DEVICE and ENABLE_UCX as CMake
 options (#3343)

No change of default value (still ON).
These were hidden cmake vars before that patch.
Fix issue #3289

Signed-off-by: William Tambellini <wtambellini@sdl.com>
Co-authored-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>
---
 cpp/CMakeLists.txt                            | 24 ++++++++++++++-----
 cpp/tensorrt_llm/kernels/CMakeLists.txt       |  2 +-
 .../kernels/cutlass_kernels/CMakeLists.txt    |  3 +--
 .../kernels/userbuffers/CMakeLists.txt        |  2 +-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 07b7889e31..428ed9c03e 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -39,6 +39,10 @@ option(FAST_BUILD "Skip compiling some kernels to accelerate compiling" OFF)
 option(FAST_MATH "Compiling in fast math mode" OFF)
 option(INDEX_RANGE_CHECK "Compiling with index range checks" OFF)
 option(COMPRESS_FATBIN "Compress everything in fatbin" ON)
+option(ENABLE_MULTI_DEVICE
+       "Enable building with multi device support (requires NCCL, MPI,...)" ON)
+option(ENABLE_UCX "Enable building with UCX (Uniform Communication X) support"
+       ON)
 
 # Always use static NVRTC for IP protection reasons.
 set(USE_SHARED_NVRTC OFF)
@@ -347,7 +351,6 @@ set(CMAKE_CUDA_RUNTIME_LIBRARY Static)
 
 find_library(RT_LIB rt)
 
-set_ifndef(ENABLE_MULTI_DEVICE 1)
 if(ENABLE_MULTI_DEVICE)
   # NCCL dependencies
   set_ifndef(NCCL_LIB_DIR /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/)
@@ -438,9 +441,15 @@ endif()
 
 # set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -G")
 
-set(CMAKE_CXX_FLAGS
-    "${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss -DENABLE_MULTI_DEVICE=${ENABLE_MULTI_DEVICE}"
-)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss ")
+
+# note: cmake expr generation $<BOOL:${ENABLE_MULTI_DEVICE}> is a build time
+# evaluation so hard to debug at cmake time
+if(ENABLE_MULTI_DEVICE)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_MULTI_DEVICE=1")
+else()
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_MULTI_DEVICE=0")
+endif()
 
 # Fix linking issue with TRT 10, the detailed description about `--mcmodel` can
 # be found in
@@ -653,7 +662,6 @@ endif()
 
 # Defer UCX/UCXX setup until after USE_CXX11_ABI is well defined, as UCXX will
 # need to be built to have aligned symbols
-set_ifndef(ENABLE_UCX 1)
 if(ENABLE_UCX)
   # Only enable UCX related features if the system has UCX library
   find_package(ucx)
@@ -680,7 +688,11 @@ if(ENABLE_UCX)
                  NO_DEFAULT_PATH)
   endif()
 endif()
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=${ENABLE_UCX}")
+if(ENABLE_UCX)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=1")
+else()
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=0")
+endif()
 
 file(STRINGS "${TRT_INCLUDE_DIR}/NvInferVersion.h" VERSION_STRINGS
      REGEX "#define NV_TENSORRT_.*")
diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt
index 4881282ccc..520ffb1b71 100644
--- a/cpp/tensorrt_llm/kernels/CMakeLists.txt
+++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt
@@ -59,7 +59,7 @@ function(filter_cuda_archs ARCH SOURCES_VAR)
   endif()
 endfunction()
 
-if(ENABLE_MULTI_DEVICE EQUAL 0)
+if(NOT ENABLE_MULTI_DEVICE)
   list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$")
 endif()
 
diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt
index 35a99a88e4..9066c59252 100644
--- a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt
+++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt
@@ -166,9 +166,8 @@ function(process_target target_name enable_hopper enable_blackwell)
       ${target_name} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-Wno-psabi>)
   endif()
 
-  target_compile_options(${target_name}
-                         PRIVATE "-DENABLE_MULTI_DEVICE=${ENABLE_MULTI_DEVICE}")
   if(ENABLE_MULTI_DEVICE)
+    target_compile_options(${target_name} PRIVATE "-DENABLE_MULTI_DEVICE=1")
     target_link_libraries(${target_name} PRIVATE ${MPI_C_LIBRARIES})
   endif()
 
diff --git a/cpp/tensorrt_llm/kernels/userbuffers/CMakeLists.txt b/cpp/tensorrt_llm/kernels/userbuffers/CMakeLists.txt
index 72082abdfe..df430c0839 100644
--- a/cpp/tensorrt_llm/kernels/userbuffers/CMakeLists.txt
+++ b/cpp/tensorrt_llm/kernels/userbuffers/CMakeLists.txt
@@ -14,7 +14,7 @@
 # License for the specific language governing permissions and limitations under
 # the License.
 #
-if(ENABLE_MULTI_DEVICE EQUAL 0)
+if(NOT ENABLE_MULTI_DEVICE)
   add_library(userbuffers_src OBJECT ub_interface.cpp userbuffersManager.cpp)
 else()
   file(GLOB_RECURSE SRC_CPP *.cpp)