                                  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);

{
    // Check if we have attained the sequence length limit. If so, stop the sequence.
    // In addition, check if all sequences are stopped and return the result in should_stop
    FT_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
    dim3 block{min(512, uint32_t(batch_size * beam_width))};
    dim3 grid{1};

    length_criterion<<<grid, block, 0, stream>>>(
        finished, should_stop, finished_sum, sequence_limit_length, batch_size, beam_width, step);
    sync_check_cuda_error();

    int h_finished_sum = 0;
    cudaD2Hcpy(&h_finished_sum, finished_sum, 1);
    *should_stop = h_finished_sum == batch_size * beam_width;
}

