Spec Scheduler

.gitignore

@@ -6,6 +6,8 @@ python/flexflow/core/flexflow_cffi_header.py
 *.pb.h
 *.o
 *.a
+*.nsys-rep
+*.nfs*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.gitmodules

@@ -22,4 +22,10 @@
 [submodule "deps/tokenizers-cpp"]
 	path = deps/tokenizers-cpp
 	url = https://github.com/mlc-ai/tokenizers-cpp.git
-	fetchRecurseSubmodules = true
+	fetchRecurseSubmodules = true
+[submodule "deps/flashinfer"]
+	path = deps/flashinfer
+	url = https://github.com/flashinfer-ai/flashinfer.git
+[submodule "deps/raft"]
+	path = deps/raft
+	url = https://github.com/rapidsai/raft.git

CMakeLists.txt

@@ -4,6 +4,12 @@ project(FlexFlow)
 
 include(ExternalProject)
 
+enable_language(CXX)
+enable_language(CUDA)
+if (CMAKE_CXX_COMPILER_VERSION VERSION_LESS 8)
+    message(FATAL_ERROR "Your C++ compiler is too old. Please upgrade to version 8 or higher.")
+endif()
+
 # Set policy CMP0074 to eliminate cmake warnings
 cmake_policy(SET CMP0074 NEW)
 cmake_policy(SET CMP0077 NEW)
@@ -128,6 +134,9 @@ list(APPEND CC_FLAGS
 list(APPEND NVCC_FLAGS
   -std=c++17)
 
+list(APPEND NVCC_FLAGS
+  --expt-relaxed-constexpr
+  --extended-lambda)
 
 add_compile_options(${CC_FLAGS})
 set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} ${NVCC_FLAGS})
@@ -201,6 +210,12 @@ if(NOT BUILD_LEGION_ONLY)
   # optional
   include(optional)
 
+  set(CMAKE_PREFIX_PATH ${CMAKE_PREFIX_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/deps/raft/cpp/build/install)
+  find_package(raft)
+  list(APPEND FLEXFLOW_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/deps/raft/cpp/include)
+
+  list(APPEND FLEXFLOW_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/deps/flashinfer/include)
+
   if (FF_GPU_BACKEND STREQUAL "cuda")
     list(APPEND FF_CC_FLAGS
       -DFF_USE_CUDA)
@@ -290,6 +305,12 @@ if(NOT BUILD_LEGION_ONLY)
       LIST_DIRECTORIES False
       ${FLEXFLOW_ROOT}/src/*.cu)
 
+    # tensorrt_llm custom allreduce
+    if(FF_USE_NCCL)
+      list(APPEND FLEXFLOW_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/deps/tensorrt_llm)
+      list(APPEND FLEXFLOW_GPU_SRC ${CMAKE_CURRENT_SOURCE_DIR}/deps/tensorrt_llm/tensorrt_llm/custom_allreduce_kernels.cu)
+    endif()
+
     add_compile_definitions(FF_USE_CUDA)
 
     if(BUILD_SHARED_LIBS)
@@ -397,6 +418,8 @@ if(NOT BUILD_LEGION_ONLY)
     target_link_libraries(flexflow ${LEGION_LIBRARY} ${FLEXFLOW_EXT_LIBRARIES} nlohmann_json::nlohmann_json mpark_variant optional)
   endif()
 
+  target_link_libraries(flexflow raft::raft)
+
   #library api version, bump from time to time
   set(SOVERSION 1)
 
@@ -425,7 +448,7 @@ if(NOT BUILD_LEGION_ONLY)
       # generate the Legion Python bindings library. When building from pip, we need to do this post-install to prevent Legion from overwriting the path to the Legion shared library
       add_custom_command(TARGET flexflow
         POST_BUILD	
-        COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/deps/legion/bindings/python/setup.py build --cmake-build-dir ${Legion_BINARY_DIR}/runtime --prefix ${Legion_BINARY_DIR} --build-lib=${Legion_BINARY_DIR}/bindings/python ${Legion_PYTHON_EXTRA_INSTALL_ARGS}
+        COMMAND CMAKE_BUILD_DIR=${Legion_BINARY_DIR}/runtime CMAKE_INSTALL_PREFIX=${Legion_BINARY_DIR} ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/deps/legion/bindings/python/setup.py build --build-lib=${Legion_BINARY_DIR}/bindings/python ${Legion_PYTHON_EXTRA_INSTALL_ARGS}
         WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/deps/legion/bindings/python
       )
       # create flexflow_python interpreter. When building from pip, we install the FF_HOME/python/flexflow_python script instead.
@@ -558,6 +581,7 @@ if(NOT BUILD_LEGION_ONLY)
   if(FF_BUILD_ALL_INFERENCE_EXAMPLES OR FF_BUILD_ALL_EXAMPLES)
     add_subdirectory(inference/spec_infer)
     add_subdirectory(inference/incr_decoding)
+    add_subdirectory(inference/trace_generator)
   endif()
 
 

FlexFlow.mk

@@ -95,9 +95,12 @@ ifneq ($(strip $(FF_USE_PYTHON)), 1)
 endif
 
 
-INC_FLAGS	+= -I${FF_HOME}/include -I${FF_HOME}/inference -I${FF_HOME}/deps/optional/include -I${FF_HOME}/deps/variant/include -I${FF_HOME}/deps/json/include -I${FF_HOME}/deps/tokenizers-cpp/include -I${FF_HOME}/deps/tokenizers-cpp/sentencepiece/src
+INC_FLAGS	+= -I${FF_HOME}/include -I${FF_HOME}/inference -I${FF_HOME}/deps/optional/include -I${FF_HOME}/deps/variant/include -I${FF_HOME}/deps/json/include -I${FF_HOME}/deps/tokenizers-cpp/include -I${FF_HOME}/deps/tokenizers-cpp/sentencepiece/src \
+				-I${FF_HOME}/deps/raft/cpp/include -I${FF_HOME}/deps/rmm/include -I${FF_HOME}/deps/spdlog/include \
+				-I${FF_HOME}/deps/flashinfer/include
 CC_FLAGS	+= -DMAX_TENSOR_DIM=$(MAX_DIM) -DLEGION_MAX_RETURN_SIZE=32768
-NVCC_FLAGS	+= -DMAX_TENSOR_DIM=$(MAX_DIM) -DLEGION_MAX_RETURN_SIZE=32768
+NVCC_FLAGS	+= -DMAX_TENSOR_DIM=$(MAX_DIM) -DLEGION_MAX_RETURN_SIZE=32768 \
+			    --expt-relaxed-constexpr --extended-lambda
 HIPCC_FLAGS     += -DMAX_TENSOR_DIM=$(MAX_DIM) -DLEGION_MAX_RETURN_SIZE=32768
 GASNET_FLAGS	+=
 # For Point and Rect typedefs

cmake/nccl.cmake

@@ -2,140 +2,88 @@ set(NCCL_NAME nccl)
 # set(NCCL_CUDA_ARCH "-gencode=arch=compute_${CUDA_ARCH},code=sm_${CUDA_ARCH}")
 # message("NCCL_CUDA_ARCH: ${NCCL_CUDA_ARCH}")
 
-set(NCCL_URL "")
-if((FF_USE_PREBUILT_NCCL OR FF_USE_ALL_PREBUILT_LIBRARIES) AND CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "x86_64")
-  if(LINUX_VERSION MATCHES "20.04")
-    if (CUDA_VERSION VERSION_EQUAL "11.0")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.0.3.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.1")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.1.1.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.2")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.2.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.3")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.3.1.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.4")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.4.3.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.5")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.5.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.6")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.6.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.7")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-20.04_11.7.0.tar.gz")
-    endif()
-  elseif(LINUX_VERSION MATCHES "18.04")
-    if (CUDA_VERSION VERSION_EQUAL "10.1")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_10.1.243.tar.gz")
-    elseif (CUDA_VERSION VERSION_EQUAL "10.2")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_10.2.89.tar.gz")
-    elseif (CUDA_VERSION VERSION_EQUAL "11.0")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.0.3.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.1")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.1.1.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.2")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.2.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.3")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.3.1.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.4")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.4.3.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.5")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.5.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.6")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.6.2.tar.gz")
-    elseif(CUDA_VERSION VERSION_EQUAL "11.7")
-      set(NCCL_URL "https://github.com/flexflow/flexflow-third-party/releases/latest/download/nccl_ubuntu-18.04_11.7.0.tar.gz")
-    endif()
-  endif()
+if(NCCL_PATH)
+  set(NCCL_ROOT ${NCCL_PATH})
+else()
+  # if NCCL_PATH is not set, let's try to find it in the CUDA root
+  set(NCCL_ROOT ${CUDA_TOOLKIT_ROOT_DIR})
 endif()
 
-if(NCCL_URL)
-  # Download and import pre-compiled NCCL library
-  message(STATUS "Using pre-compiled NCCL library")
-  message(STATUS "NCCL_URL: ${NCCL_URL}")
+find_library(NCCL_LIBRARY
+  NAMES libnccl${LIBEXT}
+  PATHS ${NCCL_ROOT} ${CUDA_ROOT}
+  PATH_SUFFIXES lib lib64
+  DOC "NCCL library." )
 
-  include(FetchContent)
-  FetchContent_Declare(${NCCL_NAME}
-    URL ${NCCL_URL}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
-  FetchContent_GetProperties(${NCCL_NAME})
-  if(NOT ${NCCL_NAME}_POPULATED)
-    FetchContent_Populate(${NCCL_NAME})
-  endif()
-
-  set(NCCL_FOLDER_PATH ${${NCCL_NAME}_SOURCE_DIR}/deps/${NCCL_NAME})
-  set(NCCL_INCLUDE_DIR ${NCCL_FOLDER_PATH}/include)
-  set(NCCL_LIB_DIR ${NCCL_FOLDER_PATH}/lib)
-  message(STATUS "NCCL library path: ${NCCL_FOLDER_PATH}")
-  add_library(nccl SHARED IMPORTED)
-  set_target_properties(nccl PROPERTIES IMPORTED_LOCATION ${NCCL_FOLDER_PATH})
+find_path(NCCL_INCLUDE_DIR
+  NAMES nccl.h
+  HINTS ${NCCL_ROOT}
+  PATH_SUFFIXES include 
+  DOC "NCCL include directory.")
 
-  list(APPEND FLEXFLOW_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
-  list(APPEND FLEXFLOW_EXT_LIBRARIES ${NCCL_LIB_DIR}/libnccl${LIBEXT})
-  install(DIRECTORY ${NCCL_INCLUDE_DIR}/ DESTINATION include)
-  install(DIRECTORY ${NCCL_LIB_DIR}/ DESTINATION lib PATTERN "pkgconfig" EXCLUDE)
-
-else()
-  if(NCCL_PATH)
-    set(NCCL_ROOT ${NCCL_PATH})
+# find NCCL, set NCCL lib and include    
+if(NCCL_LIBRARY AND NCCL_INCLUDE_DIR)
+  set(NCCL_FOUND ON)
+  set(NCCL_LIBRARIES ${NCCL_LIBRARY})
+  set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
+
+  # Check NCCL version
+  if(EXISTS "${NCCL_INCLUDE_DIR}/nccl.h")
+    file(STRINGS "${NCCL_INCLUDE_DIR}/nccl.h" NCCL_VERSION_DEFINES
+         REGEX "#define NCCL_MAJOR [0-9]+" )
+    file(STRINGS "${NCCL_INCLUDE_DIR}/nccl.h" NCCL_VERSION_DEFINES2
+         REGEX "#define NCCL_MINOR [0-9]+" )
+    string(REGEX MATCH "([0-9]+)" NCCL_MAJOR ${NCCL_VERSION_DEFINES})
+    string(REGEX MATCH "([0-9]+)" NCCL_MINOR ${NCCL_VERSION_DEFINES2})
+    set(NCCL_VERSION "${NCCL_MAJOR}.${NCCL_MINOR}")
+    if(NCCL_VERSION VERSION_LESS 2.23)
+      set(NCCL_OLD TRUE)
+    else()
+      set(NCCL_OLD FALSE)
+    endif()
+    message(STATUS "Found NCCL version: ${NCCL_VERSION}")
   else()
-    # if NCCL_PATH is not set, let's try to find it in the CUDA root
-    set(NCCL_ROOT ${CUDA_TOOLKIT_ROOT_DIR})
+    message(WARNING "NCCL header not found, unable to determine version")
+    set(NCCL_OLD TRUE)  # Assume old version if we can't determine
   endif()
-
-  find_library(NCCL_LIBRARY
-    NAMES libnccl${LIBEXT}
-    PATHS ${NCCL_ROOT} ${CUDA_ROOT}
-    PATH_SUFFIXES lib lib64
-    DOC "NCCL library." )
+endif()
 
-  find_path(NCCL_INCLUDE_DIR
-    NAMES nccl.h
-    HINTS ${NCCL_ROOT}
-    PATH_SUFFIXES include 
-    DOC "NCCL include directory.")
-
-  # find NCCL, set NCCL lib and include    
-  if(NCCL_LIBRARY AND NCCL_INCLUDE_DIR)
-    set(NCCL_FOUND ON)
-    set(NCCL_LIBRARIES ${NCCL_LIBRARY})
-    set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR})
-  endif()
-
-  # find NCCL
-  if(NCCL_FOUND)
-    list(APPEND FLEXFLOW_EXT_LIBRARIES ${NCCL_LIBRARIES})
-    list(APPEND FLEXFLOW_INCLUDE_DIRS ${NCCL_INCLUDE_DIRS})
-    message( STATUS "NCCL include : ${NCCL_INCLUDE_DIRS}" )
-    message( STATUS "NCCL libraries : ${NCCL_LIBRARIES}" )
-    add_library(nccl SHARED IMPORTED)
-
-  # Build NCCL from source
-  else()
-    message(STATUS "Building NCCL from source")
-    list(TRANSFORM CUDA_GENCODE PREPEND "NVCC_GENCODE=" OUTPUT_VARIABLE NCCL_BUILD_NVCC_GENCODE)
-
-    ExternalProject_Add(${NCCL_NAME}
-      SOURCE_DIR ${PROJECT_SOURCE_DIR}/deps/${NCCL_NAME}
-      PREFIX ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
-      INSTALL_DIR ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
-      BUILD_BYPRODUCTS ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/libnccl${LIBEXT}
-      INSTALL_COMMAND ""
-      CONFIGURE_COMMAND ""
-      BUILD_COMMAND make src.build "${NCCL_BUILD_NVCC_GENCODE}" "CUDA_HOME=${CUDA_TOOLKIT_ROOT_DIR}" "BUILDDIR=${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}"
-      BUILD_IN_SOURCE 1
-    )
+# find NCCL
+if(NCCL_FOUND AND (NOT NCCL_OLD OR CUDA_VERSION VERSION_LESS 12.0))
+  list(APPEND FLEXFLOW_EXT_LIBRARIES ${NCCL_LIBRARIES})
+  list(APPEND FLEXFLOW_INCLUDE_DIRS ${NCCL_INCLUDE_DIRS})
+  message( STATUS "NCCL include : ${NCCL_INCLUDE_DIRS}" )
+  message( STATUS "NCCL libraries : ${NCCL_LIBRARIES}" )
+  add_library(nccl SHARED IMPORTED)
+
+# Build NCCL from source
+else()
+  message(STATUS "Building NCCL from source")
+  list(TRANSFORM CUDA_GENCODE PREPEND "NVCC_GENCODE=" OUTPUT_VARIABLE NCCL_BUILD_NVCC_GENCODE)
 
-    ExternalProject_Get_Property(${NCCL_NAME} INSTALL_DIR)
-    message(STATUS "NCCL install dir: ${INSTALL_DIR}")
-    list(APPEND FLEXFLOW_INCLUDE_DIRS
-      ${INSTALL_DIR}/include)
-    list(APPEND FLEXFLOW_EXT_LIBRARIES
-      ${INSTALL_DIR}/lib/libnccl${LIBEXT})
-    set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/")
-
-    install(DIRECTORY ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/include/ DESTINATION include)
-    install(DIRECTORY ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/ DESTINATION lib PATTERN "pkgconfig" EXCLUDE)
+  set(NCCL_BUILD_CMD make src.build "${NCCL_BUILD_NVCC_GENCODE}" "CUDA_HOME=${CUDA_TOOLKIT_ROOT_DIR}" "BUILDDIR=${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}")
+  if(DEFINED ENV{MAKEFLAGS})
+    set(NCCL_BUILD_CMD ${CMAKE_COMMAND} -E env MAKEFLAGS=$ENV{MAKEFLAGS} ${NCCL_BUILD_CMD})
   endif()
+  ExternalProject_Add(${NCCL_NAME}
+    SOURCE_DIR ${PROJECT_SOURCE_DIR}/deps/${NCCL_NAME}
+    PREFIX ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
+    INSTALL_DIR ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}
+    BUILD_BYPRODUCTS ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/libnccl${LIBEXT}
+    INSTALL_COMMAND ""
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ${NCCL_BUILD_CMD}
+    BUILD_IN_SOURCE 1
+  )
 
+  ExternalProject_Get_Property(${NCCL_NAME} INSTALL_DIR)
+  message(STATUS "NCCL install dir: ${INSTALL_DIR}")
+  list(APPEND FLEXFLOW_INCLUDE_DIRS
+    ${INSTALL_DIR}/include)
+  list(APPEND FLEXFLOW_EXT_LIBRARIES
+    ${INSTALL_DIR}/lib/libnccl${LIBEXT})
+  set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/")
+
+  install(DIRECTORY ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/include/ DESTINATION include)
+  install(DIRECTORY ${CMAKE_BINARY_DIR}/deps/${NCCL_NAME}/lib/ DESTINATION lib PATTERN "pkgconfig" EXCLUDE)
 endif()

config/config.linux

@@ -111,6 +111,11 @@ function get_build_configs() {
     BUILD_CONFIGS="FF_CUDA_ARCH=${FF_CUDA_ARCH} FF_HIP_ARCH=${FF_HIP_ARCH} CUDA_DIR=${CUDA_DIR} CUDNN_DIR=${CUDNN_DIR} CUBLAS_DIR=${CUBLAS_DIR} CURAND_DIR=${CURAND_DIR} NCCL_DIR=${NCCL_DIR} FF_USE_PYTHON=${FF_USE_PYTHON} BUILD_LEGION_ONLY=${BUILD_LEGION_ONLY} FF_GASNET_CONDUIT=${FF_GASNET_CONDUIT} UCX_DIR=${UCX_DIR} FF_LEGION_NETWORKS=${FF_LEGION_NETWORKS} FF_BUILD_ALL_EXAMPLES=${FF_BUILD_ALL_EXAMPLES} FF_BUILD_ALL_INFERENCE_EXAMPLES=${FF_BUILD_ALL_INFERENCE_EXAMPLES} FF_BUILD_UNIT_TESTS=${FF_BUILD_UNIT_TESTS} FF_USE_PREBUILT_NCCL=${FF_USE_PREBUILT_NCCL} FF_USE_PREBUILT_LEGION=${FF_USE_PREBUILT_LEGION} FF_USE_ALL_PREBUILT_LIBRARIES=${FF_USE_ALL_PREBUILT_LIBRARIES} FF_USE_AVX2=${FF_USE_AVX2} FF_MAX_DIM=${FF_MAX_DIM} ROCM_PATH=${ROCM_PATH} FF_GPU_BACKEND=${FF_GPU_BACKEND} INSTALL_DIR=${INSTALL_DIR}"
 }
 
+#install raft
+echo "Building raft dependency ..."
+INSTALL_PREFIX=./install $(dirname $0)/../deps/raft/build.sh libraft > /dev/null
+echo "Building raft dependency ... Done"
+
 if [[ -n "$1" && ( "$1" == "CMAKE_FLAGS" || "$1" == "CUDA_PATH" ) ]]; then
     . $(dirname $0)/config.inc
     # Passing CMAKE_FLAGS or CUDA_PATH as $1 will print the value of the CMAKE_FLAGS/CUDA_PATH variable, 

deps/tensorrt_llm/README.md

@@ -0,0 +1,5 @@
+## Custom AllReduce Implementation
+
+This is an adapted version of the custom AllReduce plugin from NVIDIA's [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) repository.
+
+To replace the NCCL AllReduce call, we should also add a CUDA IPC support to the custom AllReduce usage. Our IPC&AllReduce implementation is referenced from [mlc-ai/relax](https://github.com/mlc-ai/relax).