 
   __syncthreads();
 
-  if (vo < pad_active_groups / 2) {
-    zero(*reinterpret_cast<V_vec *>(&out_smem[vo * Dh + vi]));
-  }
+  if (Dh == Dh_MAX || vi < Dh) {
 #pragma unroll
-  for (int active_groups = pad_active_groups; active_groups >= 2;
-       active_groups /= 2) {
-    int midpoint = active_groups / 2;
-
-    if (vo >= midpoint && vo < active_groups) {
+    for (int active_groups = V_PER_ITER; active_groups >= 2;
+         active_groups /= 2) {
+      int midpoint = active_groups / 2;
+
+      if (vo >= midpoint && vo < active_groups && (Dh == Dh_MAX || vi < Dh)) {
 #ifdef MMHA_USE_FP32_ACUM_FOR_OUT
-      convert_from_float(
-          *reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]),
-          out);
+        convert_from_float(
+            *reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]),
+            out);
 #else
-      *reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]) = out;
-#endif
+        *reinterpret_cast<V_vec *>(&out_smem[(vo - midpoint) * Dh + vi]) = out;
+#endif
+      }
+      __syncthreads();
+      if (vo < midpoint && (Dh == Dh_MAX || vi < Dh)) {
+        out =
+            add(*reinterpret_cast<const V_vec *>(&out_smem[vo * Dh + vi]), out);
+      }
+      __syncthreads();
     }
-    __syncthreads();
-    if (vo < midpoint) {
-      out = add(*reinterpret_cast<const V_vec *>(&out_smem[vo * Dh + vi]), out);
-    }
-    __syncthreads();
-  }
-
-  if (vo == 0) {
+  }
+
+  if (vo == 0 && (Dh == Dh_MAX || vi < Dh)) {
 #ifdef MMHA_USE_FP32_ACUM_FOR_OUT
     convert_from_float(*reinterpret_cast<V_vec *>(&params.out[bhi * Dh + vi]),
                        out);

  }
  __syncthreads();

  constexpr int V_VEC_SIZE = Dh / THREADS_PER_VALUE;
  using V_vec = typename V_vec_<T, V_VEC_SIZE>::Type;

  int vo = tid / THREADS_PER_VALUE;

    case 128:
      fmha_launch_kernel<T, 128, 128>(params, dev_ctx.stream());
      break;
    case 192:
      fmha_launch_kernel<T, 192, 256>(params, dev_ctx.stream());
      break;
    default:
      PADDLE_THROW(platform::errors::Unimplemented(
          "Dim_head = %d is unsupport!", dim_head));
  }
}

