                                  int             beam_width,
                                  int             step)
 {
-    typedef cub::BlockReduce<int, LENGTH_CRITERION_BLOCKSIZE> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage              temp_storage;
+    int thread_finished_count = 0;
+    for (int index = threadIdx.x; index < batch_size * beam_width; index += blockDim.x) {
+        const int batch_idx = index / beam_width;
 
-    const int index = blockIdx.x * blockDim.x + threadIdx.x;
+        finished[index] |= step >= sequence_limit_length[batch_idx];
+        thread_finished_count += finished[index] ? 1 : 0;
+    }
+    int block_finished_count = 0;
+    if (blockDim.x <= 32) {
+        block_finished_count = warpReduceSum(thread_finished_count);
+    }
+    else {
+        block_finished_count = blockReduceSum(thread_finished_count);
+    }
+    __syncthreads();
 
-    // const int beam_idx  = index % beam_width;
-    const int batch_idx = index / beam_width;
-
-    if (index >= batch_size * beam_width) {
-        return;
+    if (threadIdx.x == 0) {
+        finished_sum[0] = block_finished_count;
     }
-
-    finished[index] |= step >= sequence_limit_length[batch_idx];
-
-    int agg = BlockReduce(temp_storage).Sum((int)finished[index]);
-    atomicAdd(finished_sum, agg);
 }
 
 void invokeLengthCriterion(bool*           finished,

                              int          step,
                              cudaStream_t stream)
{
    // Check if we have sampled a word from the stop_words list. If so, stop the sequence.
    dim3 block, grid;
    block.x = min(((stop_words_len + 32 - 1) / 32) * 32, 256UL);

 * limitations under the License.
 */

#include "src/fastertransformer/kernels/stop_criteria_kernels.h"
#include "src/fastertransformer/utils/cuda_utils.h"
#include "src/fastertransformer/utils/memory_utils.h"
#include "src/fastertransformer/kernels/reduce_kernel_utils.cuh"

namespace fastertransformer {

__global__ void stop_words_criterion(const int* output_ids,
                                     const int* parent_ids,
