comm: refactor and initialize flashinfer.comm module (#1089)

Author: yzh119

.gitmodules

@@ -4,9 +4,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
-[submodule "3rdparty/mscclpp"]
-	path = 3rdparty/mscclpp
-	url = https://github.com/microsoft/mscclpp.git
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
 	url = https://github.com/NVIDIA/cutlass.git

3rdparty/mscclpp

@@ -49,9 +49,6 @@ flashinfer_option(
 flashinfer_option(
   FLASHINFER_NORM
   "Whether to compile normalization kernel tests/benchmarks or not." OFF)
-flashinfer_option(
-  FLASHINFER_DISTRIBUTED
-  "Whether to compile distributed kernel tests/benchmarks or not." OFF)
 flashinfer_option(FLASHINFER_FASTDIV_TEST
                   "Whether to compile fastdiv kernel tests or not." OFF)
 flashinfer_option(FLASHINFER_FASTDEQAUNT_TEST
@@ -83,11 +80,7 @@ if(FLASHINFER_PREFILL
    OR FLASHINFER_NORM)
   message(STATUS "NVBench and GoogleTest enabled")
   add_subdirectory(3rdparty/nvbench)
-  if(FLASHINFER_DISTRIBUTED)
-    add_subdirectory(3rdparty/mscclpp)
-  else(FLASHINFER_DISTRIBUTED)
-    add_subdirectory(3rdparty/googletest)
-  endif(FLASHINFER_DISTRIBUTED)
+  add_subdirectory(3rdparty/googletest)
 endif(
   FLASHINFER_PREFILL
   OR FLASHINFER_DECODE
@@ -457,29 +450,3 @@ if(FLASHINFER_FASTDEQUANT_TEST)
     test_fast_dequant PRIVATE ${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR})
   target_link_libraries(test_fast_dequant PRIVATE gtest gtest_main)
 endif(FLASHINFER_FASTDEQUANT_TEST)
-
-if(FLASHINFER_DISTRIBUTED)
-  find_package(MPI REQUIRED)
-
-  message(STATUS "Compile sum all-reduce kernel tests.")
-  file(GLOB_RECURSE TEST_DIST_SUM_ALL_REDUCE_SRCS
-       ${PROJECT_SOURCE_DIR}/src/test_sum_all_reduce.cu)
-  add_executable(test_sum_all_reduce ${TEST_DIST_SUM_ALL_REDUCE_SRCS})
-  target_include_directories(
-    test_sum_all_reduce
-    PRIVATE ${FLASHINFER_INCLUDE_DIR} 3rdparty/mscclpp/include
-            3rdparty/spdlog/include)
-  target_link_libraries(test_sum_all_reduce PRIVATE MPI::MPI_CXX mscclpp)
-  target_compile_definitions(test_sum_all_reduce PRIVATE -DENABLE_MPI)
-
-  message(STATUS "Compile attention allreduce kernel tests.")
-  file(GLOB_RECURSE TEST_DIST_ATTN_ALL_REDUCE_SRCS
-       ${PROJECT_SOURCE_DIR}/src/test_attn_all_reduce.cu)
-  add_executable(test_attn_all_reduce ${TEST_DIST_ATTN_ALL_REDUCE_SRCS})
-  target_include_directories(
-    test_attn_all_reduce
-    PRIVATE ${FLASHINFER_INCLUDE_DIR} 3rdparty/mscclpp/include
-            3rdparty/spdlog/include)
-  target_link_libraries(test_attn_all_reduce PRIVATE MPI::MPI_CXX mscclpp)
-  target_compile_definitions(test_attn_all_reduce PRIVATE -DENABLE_MPI)
-endif(FLASHINFER_DISTRIBUTED)

CMakeLists.txt

@@ -19,8 +19,6 @@ set(FLASHINFER_NORM ON)
 set(FLASHINFER_FASTDIV_TEST ON)
 # Whether to compile fastdequant tests
 set(FLASHINFER_FASTDEQUANT_TEST ON)
-# Whether to compile distributed tests
-set(FLASHINFER_DISTRIBUTED ON)
 # The following configurations can impact the binary size of the generated
 # library
 set(FLASHINFER_GEN_HEAD_DIMS 64 128 256 512)

cmake/config.cmake

@@ -1,6 +1,6 @@
 // flashinfer: adapted from sglang + vllm code
 // refer to: https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
-#include "flashinfer/distributed/custom_all_reduce.cuh"
+#include "flashinfer/comm/custom_all_reduce.cuh"
 #include "pytorch_extension_utils.h"
 
 // Fake pointer type, must match fptr_t type in ops.h.

csrc/custom_all_reduce.cu

@@ -10,7 +10,7 @@
 
 from .activation import act_func_def_str, gen_act_and_mul_module
 from .cascade import gen_cascade_module
-from .custom_all_reduce import gen_comm_module
+from .comm import gen_comm_module
 from .gemm import gen_gemm_module, gen_gemm_sm90_module, gen_gemm_sm100_module
 from .jit import JitSpec, build_jit_specs
 from .jit import env as jit_env