The code assumes WARP_SIZE to be equal to 32, which is not the case o…

…n ROCm (#406) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
ROCm · Feb 5, 2025 · f65ecc9 · f65ecc9
1 parent ed3337d
commit f65ecc9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu
@@ -207,8 +207,8 @@ __global__ void sgl_moe_align_block_size_kernel(
   __shared__ int32_t shared_counts[32][8];
   __shared__ int32_t local_offsets[256];
 
-  const int warp_id = threadIdx.x / WARP_SIZE;
-  const int lane_id = threadIdx.x % WARP_SIZE;
+  const int warp_id = threadIdx.x / 32;
+  const int lane_id = threadIdx.x % 32;
   const int experts_per_warp = 8;
   const int my_expert_start = warp_id * experts_per_warp;