[DCU] Add width to __shfl

Signed-off-by: wenjh <wenjh@sugon.com>

[DCU] Add width to __shfl
Signed-off-by: wenjh <wenjh@sugon.com>
6ed9a3e4 · wenjh · b27e513d · 6ed9a3e4
Commit 6ed9a3e4 authored May 22, 2025 by wenjh
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu ...e/common/transpose/quantize_transpose_square_blockwise.cu +4 -2

No files found.
--- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu
@@ -135,7 +135,8 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
  // broadcast the amax to all threads in a warp from the lane 0
  constexpr int lane_zero = 0;
 #ifdef __HIP_PLATFORM_AMD__
-  warp_tile_amax = __shfl(warp_tile_amax, lane_zero);
+  warp_tile_amax = __shfl(warp_tile_amax, lane_zero, THREADS_PER_WARP);
+  __syncthreads();
 #else
  warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);
 #endif
@@ -362,7 +363,8 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose
  // broadcast the amax to all threads in a warp from the lane 0
  constexpr int lane_zero = 0;
 #ifdef __HIP_PLATFORM_AMD__
-  warp_tile_amax = __shfl(warp_tile_amax, lane_zero);
+  warp_tile_amax = __shfl(warp_tile_amax, lane_zero, THREADS_PER_WARP);
+  __syncthreads();
 #else
  warp_tile_amax = __shfl_sync(0xFFFFFFFF, warp_tile_amax, lane_zero);
 #endif