add FlagCX as optional communication library when using iluvatar gpus (#1862)

Author: mikethegoblin

backends/iluvatar_gpu/CMakeLists.txt

@@ -33,6 +33,10 @@ include(external/eigen)
 include(external/xxhash)
 include(external/zlib)
 include(external/protobuf)
+if(WITH_FLAGCX)
+  add_definitions("-DPADDLE_WITH_FLAGCX")
+  include(external/flagcx)
+endif()
 
 set(PLUGIN_VERSION ${PADDLE_VERSION})
 set(PROTO_FILE "${PADDLE_SOURCE_DIR}/paddle/phi/core/external_error.proto")
@@ -66,8 +70,11 @@ target_include_directories(external_error_proto
 target_link_libraries(external_error_proto PUBLIC protobuf)
 set_target_properties(external_error_proto PROPERTIES POSITION_INDEPENDENT_CODE
                                                       ON)
-
-add_custom_target(external_deps DEPENDS eigen3 zlib protobuf)
+if(WITH_FLAGCX)
+  add_custom_target(external_deps DEPENDS eigen3 zlib protobuf flagcx)
+else()
+  add_custom_target(external_deps DEPENDS eigen3 zlib protobuf)
+endif()
 
 if(WITH_COREX)
   add_definitions(-DPADDLE_WITH_COREX)
@@ -247,6 +254,10 @@ file(
   kernels/ernie_core/*.cc
   kernels/gpudnn/*.cu)
 
+if(WITH_FLAGCX)
+  list(APPEND CC_SRCS runtime/runtime_flagcx.cc)
+endif()
+
 set(CUSTOM_DEVICE_SRCS ${CUDA_SRCS} ${CC_SRCS})
 
 set_source_files_properties(${CUSTOM_DEVICE_SRCS} PROPERTIES LANGUAGE CUDA)
@@ -269,7 +280,9 @@ target_link_libraries(
           protobuf
           external_error_proto
           cuinfer
-          nccl)
+          nccl
+          # change nccl to ${FLAGCX_LIB} if compiling with FlagCX ${FLAGCX_LIB}
+)
 
 include_directories(BEFORE ${PADDLE_SOURCE_DIR})
 

backends/iluvatar_gpu/build_paddle.sh

@@ -1,13 +1,13 @@
 #!/bin/bash
 
 # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-# 
+#
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
-# 
+#
 #     http://www.apache.org/licenses/LICENSE-2.0
-# 
+#
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -26,6 +26,15 @@ export CMAKE_CUDA_ARCHITECTURES=${COREX_ARCH}
 CURRENT_DIR=$(pwd)
 PADDLE_SOURCE_DIR="${CURRENT_DIR}/../../Paddle"
 PATCH_FILE="${CURRENT_DIR}/patches/paddle-corex.patch"
+# set BUILD_WITH_FLAGCX to 1 if we want to use flagcx as communication backend
+BUILD_WITH_FLAGCX=0
+FLAGCX_ROOT="/workspace/FlagCX"
+
+if [ "$BUILD_WITH_FLAGCX" == "1" ]; then
+    WITH_FLAGCX="ON"
+else
+    WITH_FLAGCX="OFF"
+fi
 
 bash clean_paddle.sh
 
@@ -51,9 +60,10 @@ if [[ ! -d "build" ]]; then
 fi
 pushd build
 
-cmake -DPY_VERSION=${PYTHON_VERSION} -DWITH_COREX=ON \
--DWITH_DISTRIBUTE=ON -DWITH_NCCL=ON -DWITH_RCCL=OFF -DCMAKE_BUILD_TYPE=Release \
+cmake -DPY_VERSION=${PYTHON_VERSION} -DWITH_COREX=ON -DPADDLE_SOURCE_DIR=${PADDLE_SOURCE_DIR} \
+-DWITH_DISTRIBUTE=ON -DWITH_NCCL=ON -DWITH_FLAGCX=${WITH_FLAGCX} -DWITH_RCCL=OFF -DCMAKE_BUILD_TYPE=Release \
 -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DON_INFER=ON -DCOREX_VERSION=${COREX_VERSION} -DCOREX_ARCH=${COREX_ARCH} \
+-DFLAGCX_ROOT=${FLAGCX_ROOT} \
 -DCMAKE_CXX_FLAGS='-Wno-error=pessimizing-move -Wno-error=deprecated-copy -Wno-error=init-list-lifetime' \
 -DCMAKE_CUDA_FLAGS='-Xclang -fcuda-allow-variadic-functions -mllvm --skip-double' \
 -DWITH_ARM=OFF -DWITH_DGC=OFF .. 2>&1 | tee compile.log

backends/iluvatar_gpu/cmake/external/flagcx.cmake

@@ -0,0 +1,41 @@
+set(CMAKE_FIND_DEBUG_MODE ON)
+# flagcx.cmake
+if(NOT WITH_FLAGCX)
+  return()
+endif()
+
+set(FLAGCX_SOURCE_DIR "${FLAGCX_ROOT}")
+set(FLAGCX_LIB_DIR "${FLAGCX_SOURCE_DIR}/build/lib")
+set(FLAGCX_BINARY_DIR "${PADDLE_SOURCE_DIR}/build/third_party/flagcx")
+set(THIRD_PARTY_DIR "${PADDLE_SOURCE_DIR}/build/third_party")
+
+file(REMOVE_RECURSE ${FLAGCX_BINARY_DIR})
+message(STATUS "removed old flagcx dir")
+message(STATUS "Copying third-party source to build directory")
+execute_process(COMMAND cp -r ${FLAGCX_SOURCE_DIR} ${THIRD_PARTY_DIR}
+                RESULT_VARIABLE COPY_RESULT)
+
+if(NOT COPY_RESULT EQUAL 0)
+  message(FATAL_ERROR "Failed to copy third-party source to build directory")
+endif()
+
+# Create a custom target to build the third-party library
+message(STATUS "Building third-party library with its Makefile")
+
+find_path(
+  FLAGCX_INCLUDE_DIR flagcx.h
+  PATHS ${FLAGCX_SOURCE_DIR}/flagcx/include
+  NO_DEFAULT_PATH)
+
+message(STATUS "FLAGCX_INCLUDE_DIR is ${FLAGCX_INCLUDE_DIR}")
+include_directories(SYSTEM ${FLAGCX_INCLUDE_DIR})
+
+add_library(flagcx INTERFACE)
+find_library(
+  FLAGCX_LIB
+  NAMES flagcx libflagcx
+  PATHS ${FLAGCX_LIB_DIR}
+  DOC "My custom library")
+
+add_dependencies(flagcx FLAGCX_LIB)
+message(STATUS "FLAGCX_LIB is ${FLAGCX_LIB}")

backends/iluvatar_gpu/runtime/runtime.cc

@@ -15,6 +15,9 @@
 #include <errno.h>
 #include <fcntl.h>
 #include <nccl.h>
+#if defined(PADDLE_WITH_FLAGCX)
+#include "runtime_flagcx.h"  // NOLINT
+#endif
 #include <semaphore.h>
 #include <sys/types.h>
 #include <sys/wait.h>
@@ -1023,6 +1026,9 @@ void InitPlugin(CustomRuntimeParams *params) {
          0,
          sizeof(C_DeviceInterface));
 
+#if defined(PADDLE_WITH_FLAGCX)
+  flagcxHandleInit(&flagcx_handler);
+#endif
   params->interface->get_compute_capability = GetComputeCapability;
   params->interface->get_runtime_version = GetRuntimeVersion;
   params->interface->get_driver_version = GetDriverVersion;
@@ -1073,6 +1079,22 @@ void InitPlugin(CustomRuntimeParams *params) {
   params->interface->init_eigen_device = InitEigenDevice;
   params->interface->destroy_eigen_device = DestroyEigenDevice;
 
+#if defined(PADDLE_WITH_FLAGCX)
+  params->interface->xccl_all_gather = XcclFlagcxAllGather;
+  params->interface->xccl_all_reduce = XcclFlagcxAllReduce;
+  params->interface->xccl_broadcast = XcclFlagcxBroadcast;
+  params->interface->xccl_comm_init_rank = XcclFlagcxCommInitRank;
+  params->interface->xccl_destroy_comm = XcclFlagcxDestroyComm;
+  params->interface->xccl_get_unique_id = XcclFlagcxGetUniqueId;
+  params->interface->xccl_get_unique_id_size = XcclFlagcxGetUniqueIdSize;
+  params->interface->xccl_group_end = XcclFlagcxGroupEnd;
+  params->interface->xccl_group_start = XcclFlagcxGroupStart;
+  params->interface->xccl_recv = XcclFlagcxRecv;
+  params->interface->xccl_reduce = XcclFlagcxReduce;
+  params->interface->xccl_reduce_scatter = XcclFlagcxReduceScatter;
+  params->interface->xccl_send = XcclFlagcxSend;
+  params->interface->xccl_all_to_all = XcclFlagcxAllToAll;
+#else
   params->interface->xccl_all_gather = XcclAllGather;
   params->interface->xccl_all_reduce = XcclAllReduce;
   params->interface->xccl_broadcast = XcclBroadcast;
@@ -1086,6 +1108,7 @@ void InitPlugin(CustomRuntimeParams *params) {
   params->interface->xccl_reduce = XcclReduce;
   params->interface->xccl_reduce_scatter = XcclReduceScatter;
   params->interface->xccl_send = XcclSend;
+#endif
 
   params->interface->profiler_collect_trace_data = nullptr;
   params->interface->profiler_initialize = nullptr;