From 1f6dab4b34cdfe83eddcba9ed90468dd38845bf2 Mon Sep 17 00:00:00 2001
From: fruitea <aetiurf@gmail.com>
Date: Sat, 23 Nov 2024 01:04:15 -0800
Subject: [PATCH] fix: AllReduce minor

---
 src/parallel_ops/kernels/allreduce_kernels.cu | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/parallel_ops/kernels/allreduce_kernels.cu b/src/parallel_ops/kernels/allreduce_kernels.cu
index 2dc1caf19f..0e5c15008e 100644
--- a/src/parallel_ops/kernels/allreduce_kernels.cu
+++ b/src/parallel_ops/kernels/allreduce_kernels.cu
@@ -142,6 +142,9 @@ void inference_kernel_wrapper(Context ctx,
   int device_id = m->handle.device_id;
   ncclComm_t ncclComm = m->handle.ncclComm;
   DataType dtype = input.data_type;
+  if (num_elements == 0) {
+    return;
+  }
 
   tensorrt_llm::AllReduceStrategyType strategy =
       tensorrt_llm::SelectImplementation(