  * 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          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;
}

