Fix beam search acc issue in workgroup reduce (#3771) (#3796)

* Fix beam search acc issue in workgroup reduce (#3771)

Signed-off-by: majing <[email protected]>

Author: majing921201

csrc/gpu/aten/operators/BeamSearch.h

@@ -201,6 +201,8 @@ void beam_search_topk_stage2(
     int64_t batch_size,
     int32_t num_wg_per_beam) {
   int slm_size = (sizeof(int) + sizeof(scalar_t)) * MAX_K * num_wg_per_beam;
+  int sub_wg_size = dpcppMaxSubGroupSize();
+  int wg_size = (num_wg_per_beam + sub_wg_size - 1) / sub_wg_size * sub_wg_size;
   auto& dpcpp_queue = dpcppGetCurrentQueue();
   auto cgf = DPCPP_Q_CGF(cgh) {
     dpcpp_local_acc_t<char> shared(slm_size, cgh);
@@ -233,7 +235,7 @@ void beam_search_topk_stage2(
           TopK<scalar_t, MAX_K>,
           decltype(item),
           decltype(combine),
-          1>(item, num_wg_per_beam, shared, value, combine);
+          1>(item, item.get_local_range(0), shared, value, combine);
       TopK<scalar_t, MAX_K> total = value[0];
 
       if (wi_id == 0) {
@@ -249,8 +251,8 @@ void beam_search_topk_stage2(
 
     cgh.parallel_for(
         sycl::nd_range<1>(
-            sycl::range<1>(batch_size * beam_size * num_wg_per_beam),
-            sycl::range<1>(num_wg_per_beam)),
+            sycl::range<1>(batch_size * beam_size * wg_size),
+            sycl::range<1>(wg_size)),
         kfn);
   };
   DPCPP_Q_SUBMIT(dpcpp_queue, cgf);
@@ -667,7 +669,9 @@ void finalize(
     const int64_t out_sentence_num,
     const int64_t pad_token_id) {
   auto& dpcpp_queue = dpcppGetCurrentQueue();
-  int32_t wg_size = 2 * beam_size;
+  int32_t sub_wg_size = dpcppMaxSubGroupSize();
+  int32_t wg_size =
+      (2 * beam_size + sub_wg_size - 1) / sub_wg_size * sub_wg_size;
 
   auto cgf = DPCPP_Q_CGF(cgh) {
     auto s_score = dpcpp_local_acc_t<scalar_t>(wg_size, cgh);
@@ -676,7 +680,7 @@ void finalize(
       const int32_t wi_id = item.get_local_id(0);
       const int32_t wg_id = item.get_group(0);
       if (wi_id < beam_hyps_num_beams[wg_id]) {
-        s_score[wi_id] = beam_hyps_normed_scores[wg_id * wg_size + wi_id];
+        s_score[wi_id] = beam_hyps_normed_scores[wg_id * 2 * beam_size + wi_id];
       } else {
         s_score[wi_id] = std::numeric_limits<scalar_t>::lowest();
       }

csrc/gpu/aten/operators/Reduce.h

@@ -41,6 +41,9 @@ inline at::detail::Array<arg_t, out_vec_sz> group_reduce(
   int sg_lid = sg.get_local_linear_id();
   int sg_gid = sg.get_group_linear_id();
   int sg_range = sg.get_group_range()[0];
+  // group reduce requests workgroup size is multiple of subgroup size
+  SYCL_KERNEL_ASSERT(
+      wg_size % sg_size == 0 && "unsupported workgroup size for group reduce");
 
   for (int offset = 1; offset < sg_size; offset <<= 1) {
 #pragma unroll(out_vec_sz)