[AMD][ROCm] Fixes review comments

part  2

Signed-off-by: Artem Kuzmitckii <Artem Kuzmitckii artem.kuzmitckii@amd.com>

Author: Artem Kuzmitckii

csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h

@@ -233,11 +233,7 @@ struct call_conditional<false, TA, TB> {
 
 CUTLASS_DEVICE int32_t warp_uniform(int32_t value)
 {
-#if defined(__HIP_PLATFORM_AMD__)
-    return (int32_t)__shfl_sync(static_cast<uint64_t>(0xffffffff), (unsigned)value, 0);
-#else
     return (int32_t)__shfl_sync(0xffffffff, (unsigned)value, 0);
-#endif
 }
 
 template <typename T>

csrc/includes/reduction_utils.h

@@ -558,7 +558,7 @@ DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
 #pragma unroll
     for (int i = 1; i < reduce_width; i *= 2) {
         data[0] = element<Op1>(data[0], shfl_xor_helper(warp, data[0], i));
-        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[0], i));
+        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[1], i));
     }
 }
 
@@ -568,8 +568,8 @@ DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
 #pragma unroll
     for (int i = 1; i < reduce_width; i *= 2) {
         data[0] = element<Op1>(data[0], shfl_xor_helper(warp, data[0], i));
-        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[0], i));
-        data[2] = element<Op3>(data[2], shfl_xor_helper(warp, data[0], i));
+        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[1], i));
+        data[2] = element<Op3>(data[2], shfl_xor_helper(warp, data[2], i));
     }
 }
 
@@ -584,9 +584,9 @@ DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
 #pragma unroll
     for (int i = 1; i < reduce_width; i *= 2) {
         data[0] = element<Op1>(data[0], shfl_xor_helper(warp, data[0], i));
-        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[0], i));
-        data[2] = element<Op3>(data[2], shfl_xor_helper(warp, data[0], i));
-        data[3] = element<Op4>(data[3], shfl_xor_helper(warp, data[0], i));
+        data[1] = element<Op2>(data[1], shfl_xor_helper(warp, data[1], i));
+        data[2] = element<Op3>(data[2], shfl_xor_helper(warp, data[2], i));
+        data[3] = element<Op4>(data[3], shfl_xor_helper(warp, data[3], i));
     }
 }
 

deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/utils_paralleldequant.cuh

@@ -120,11 +120,7 @@ __device__ __forceinline__ void ExtractFromSharedToReg_Scales(uint32_t* Scales,
 #pragma unroll
     for (int i = 0; i < 4; i++) {
         // T __shfl_sync(unsigned mask, T var, int srcLane, int width=warpSize);
-#if defined(__HIP_PLATFORM_AMD__)
-        Scales[i] = __shfl_sync(static_cast<uint64_t>(0xffffffff), tmpReg, i, 4);
-#else
         Scales[i] = __shfl_sync(0xffffffff, tmpReg, i, 4);
-#endif
     }
 }