  * limitations under the License.
  */
 
-#ifndef CUDART_VERSION
-#error CUDART_VERSION Undefined!
-#elif (CUDART_VERSION >= 11050)
-#include <cub/cub.cuh>
-#else
-#include "3rdparty/cub/cub.cuh"
-#endif
-
 #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 {
-
-constexpr int LENGTH_CRITERION_BLOCKSIZE = 256;
 
 __global__ void stop_words_criterion(const int* output_ids,
                                      const int* parent_ids,

                                 int             beam_width,
                                 int             step)
{
    typedef cub::BlockReduce<int, LENGTH_CRITERION_BLOCKSIZE> BlockReduce;
    __shared__ typename BlockReduce::TempStorage              temp_storage;

    const int index = blockIdx.x * blockDim.x + threadIdx.x;

    // const int beam_idx  = index % beam_width;
    const int batch_idx = index / beam_width;

    if (index >= batch_size * beam_width) {
        return;
    }

    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)
{
    FT_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
    // 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);
