Unify two versions of AllReduce custom op (NVIDIA#3032)

* Rewrite unit test for unified allreduce op. Removing the legacy unit test.
* Revise formats, fusion_op bindings. Put all tensors as optional inputs.
* Move the MoeAllreduceOp to a separate custom op.
* Move all the fusion patterns to the new version of the AllReduce fusion kernel. Remove the AllReduce strategy config. Revise the AllReduce strategies and fusion pattern definitions.
* Add more TODOs, fixing minor bugs, and remove legacy code.

Signed-off-by: Yukun He <[email protected]>

Author: hyukn

cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu

@@ -615,7 +615,7 @@ void allreduce_fusion_kernel_launcher(AllReduceFusionParams const& params)
     TLLM_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
     static int SM = tensorrt_llm::common::getSMVersion();
     int token_num = params.size / params.hidden_dim;
-    bool oneshot = use_oneshot(token_num);
+    bool oneshot = params.use_oneshot;
     int cluster_num = token_num;
     std::array<int, NRanks> begin_tokens, token_num_per_ranks;
     if (!oneshot)

cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h

@@ -131,6 +131,7 @@ struct AllReduceFusionParams
     void* rms_gamma;
     float rms_eps;
     float* scale_factor;
+    bool use_oneshot;
     FP4QuantizationSFLayout layout = FP4QuantizationSFLayout::SWIZZLED;
     cudaStream_t stream;
     AllReduceFusionPattern pattern;

cpp/tensorrt_llm/kernels/customAllReduceKernels.h

@@ -46,10 +46,11 @@ static constexpr int kLamportHiddenSizeThreshold = 256;
 enum class AllReduceStrategyType : int8_t
 {
     NCCL = 0,
-    ONESHOT = 1,
-    TWOSHOT = 2,
-    UB = 3,
-    AUTO = 4,
+    MIN_LATENCY = 1,
+    UB = 2,
+    AUTO = 3,
+    ONESHOT = 4,
+    TWOSHOT = 5,
 };
 
 enum class AllReduceStrategyConfig : int8_t
@@ -66,10 +67,36 @@ enum class AllReduceFusionOp : int8_t
     RESIDUAL_RMS_PREPOST_NORM = 3,
     RESIDUAL_RMS_NORM_QUANT_FP8 = 4,
     RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5,
-    MOE_ALLREDUCE_RESIDUAL_RMS_NORM = 6,
-    RESIDUAL_RMS_NORM_AND_QUANT_NVFP4 = 7,
+    RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6,
+    RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7,
+    MOE_ALLREDUCE_RESIDUAL_RMS_NORM = 8,
 };
 
+inline std::ostream& operator<<(std::ostream& os, AllReduceFusionOp op)
+{
+    switch (op)
+    {
+    case AllReduceFusionOp::NONE: os << "NONE"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_NORM: os << "RESIDUAL_RMS_NORM"; break;
+    case AllReduceFusionOp::LAST_PROCESS_FOR_UB: os << "LAST_PROCESS_FOR_UB"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM: os << "RESIDUAL_RMS_PREPOST_NORM"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8: os << "RESIDUAL_RMS_NORM_QUANT_FP8"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4: os << "RESIDUAL_RMS_NORM_QUANT_NVFP4"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8: os << "RESIDUAL_RMS_NORM_OUT_QUANT_FP8"; break;
+    case AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4: os << "RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4"; break;
+    case AllReduceFusionOp::MOE_ALLREDUCE_RESIDUAL_RMS_NORM: os << "MOE_ALLREDUCE_RESIDUAL_RMS_NORM"; break;
+    default: os << "UNKNOWN"; break;
+    }
+    return os;
+}
+
+inline std::string toString(AllReduceFusionOp op)
+{
+    std::ostringstream oss;
+    oss << op;
+    return oss.str();
+}
+
 struct AllReduceFusionParams
 {
     AllReduceFusionParams()

cpp/tensorrt_llm/pybind/runtime/bindings.cpp

@@ -16,7 +16,9 @@
  */
 
 #include "bindings.h"
+#include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h"
 #include "tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.h"
+#include "tensorrt_llm/kernels/customAllReduceKernels.h"
 #include "tensorrt_llm/kernels/delayStream.h"
 #include "tensorrt_llm/runtime/cudaStream.h"
 #include "tensorrt_llm/runtime/decodingInput.h"
@@ -413,6 +415,26 @@ void initBindings(pybind11::module_& m)
             tensorrt_llm::kernels::invokeDelayStreamKernel(delay_micro_secs, stream);
         },
         "Delay kernel launch on the default stream");
+
+    py::enum_<tensorrt_llm::kernels::AllReduceFusionOp>(m, "AllReduceFusionOp")
+        .value("NONE", tensorrt_llm::kernels::AllReduceFusionOp::NONE)
+        .value("RESIDUAL_RMS_NORM", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM)
+        .value("LAST_PROCESS_FOR_UB", tensorrt_llm::kernels::AllReduceFusionOp::LAST_PROCESS_FOR_UB)
+        .value("RESIDUAL_RMS_PREPOST_NORM", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM)
+        .value("RESIDUAL_RMS_NORM_QUANT_FP8", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8)
+        .value("RESIDUAL_RMS_NORM_QUANT_NVFP4", tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4)
+        .value("RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4",
+            tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4)
+        .value("RESIDUAL_RMS_NORM_OUT_QUANT_FP8",
+            tensorrt_llm::kernels::AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8);
+
+    py::enum_<tensorrt_llm::kernels::AllReduceStrategyType>(m, "AllReduceStrategy")
+        .value("NCCL", tensorrt_llm::kernels::AllReduceStrategyType::NCCL)
+        .value("MIN_LATENCY", tensorrt_llm::kernels::AllReduceStrategyType::MIN_LATENCY)
+        .value("AUTO", tensorrt_llm::kernels::AllReduceStrategyType::AUTO)
+        .value("UB", tensorrt_llm::kernels::AllReduceStrategyType::UB)
+        .value("ONESHOT", tensorrt_llm::kernels::AllReduceStrategyType::ONESHOT)
+        .value("TWOSHOT", tensorrt_llm::kernels::AllReduceStrategyType::TWOSHOT);
 }
 
 } // namespace tensorrt_llm::pybind::runtime