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 <[email protected]>
Co-authored-by: Yuan Tong <[email protected]>

Author: WilliamTambellini

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_.*")

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()
 

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()
 

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)