Fix warp reduce issue.

Signed-off-by: Shiyu Li <[email protected]>

Author: timlee0212

cpp/tensorrt_llm/kernels/communicationKernels/mnnvlTwoShotAllreduceKernels.cu

@@ -377,11 +377,14 @@ inline __device__ T add(T a, T b)
 
 #define FINAL_MASK 0xffffffff
 #define WARP_SIZE 32
+#define LOG2_WARP_SIZE 5
+#define LANE_ID_MASK 0x1f
 
 template <typename T>
-__inline__ __device__ T warpReduceSum(T val)
+__inline__ __device__ T warpReduceSumPartial(T val)
 {
-    int lane_id = threadIdx.x & 0x1f;
+    int lane_id = threadIdx.x & LANE_ID_MASK;
+    // This function should only be called on the last warp
     int warp_size = blockDim.x - (threadIdx.x & ~(WARP_SIZE - 1));
     unsigned int active_mask = (1U << warp_size) - 1;
 
@@ -395,14 +398,25 @@ __inline__ __device__ T warpReduceSum(T val)
     return val;
 }
 
+template <typename T>
+inline __device__ T warpReduceSumFull(T val)
+{
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+    {
+        val = add<T>(val, __shfl_xor_sync(FINAL_MASK, val, offset, WARP_SIZE));
+    }
+    return val;
+}
 inline __device__ float block_reduce_sum(float val)
 {
     __shared__ float smem[WARP_SIZE];
-    int lane_id = threadIdx.x % WARP_SIZE;
-    int warp_id = threadIdx.x / WARP_SIZE;
-    int warp_num = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE; // Ceiling division to include partial warps
+    int lane_id = threadIdx.x & LANE_ID_MASK;
+    int warp_id = threadIdx.x >> LOG2_WARP_SIZE;
+    bool has_partial_warp = (blockDim.x % WARP_SIZE) != 0;
+    int warp_num = (blockDim.x + WARP_SIZE - 1) >> LOG2_WARP_SIZE; // Ceiling division to include partial warps
 
-    val = warpReduceSum(val);
+    val = (has_partial_warp && (warp_id == warp_num - 1)) ? warpReduceSumPartial(val) : warpReduceSumFull(val);
     if (lane_id == 0)
     {
         smem[warp_id] = val;
@@ -411,10 +425,10 @@ inline __device__ float block_reduce_sum(float val)
     if (warp_id == 0)
     {
         val = lane_id < warp_num ? smem[lane_id] : 0.f;
-        val = warpReduceSum(val);
+        val = (has_partial_warp && warp_num == 1) ? warpReduceSumPartial(val) : warpReduceSumFull(val);
         if (lane_id == 0)
         {
-            val = smem[0];
+            smem[0] = val;
         }
     }
     __syncthreads();