// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/phi/kernels/top_p_sampling_kernel.h"

#ifdef PADDLE_WITH_HIP
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#include <hiprand_kernel.h>
#else
#include <cuda_fp16.h>
#include <curand_kernel.h>
#endif

#include "paddle/phi/kernels/funcs/cub.h"
#if defined(__CUDACC__) && CUDA_VERSION >= 11060
#define CUDA_BFLOAT16_AVAILABLE
#include <cuda_bf16.h>
#endif

#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_device_function.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/funcs/gather.cu.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/funcs/top_k_function_cuda.h"
#include "paddle/phi/kernels/primitive/functor_primitives.h"

#ifdef PADDLE_WITH_HIP
#define GPU(str) hip##str
#else
#define GPU(str) cu##str
#endif

// #define DEBUG_TOPP

namespace phi {

template <typename T>
struct DataTypeTraits {
  using DataType = T;
};

template <>
struct DataTypeTraits<phi::float16> {
  using DataType = half;
};

#ifdef CUDA_BFLOAT16_AVAILABLE
template <>
struct DataTypeTraits<phi::bfloat16> {
  using DataType = __nv_bfloat16;
};
#endif

#define FINAL_MASK 0xFFFFFFFF

#define FIXED_BLOCK_DIM_BASE(dim, ...) \
  case (dim): {                        \
    constexpr auto kBlockDim = (dim);  \
    __VA_ARGS__;                       \
  } break

#ifdef PADDLE_WITH_HIP
#define WARP_SIZE 64
#define FIXED_BLOCK_DIM(...)                 \
  FIXED_BLOCK_DIM_BASE(1024, ##__VA_ARGS__); \
  FIXED_BLOCK_DIM_BASE(512, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(256, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(128, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(64, ##__VA_ARGS__);
#else
#define WARP_SIZE 32
#define FIXED_BLOCK_DIM(...)                 \
  FIXED_BLOCK_DIM_BASE(1024, ##__VA_ARGS__); \
  FIXED_BLOCK_DIM_BASE(512, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(256, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(128, ##__VA_ARGS__);  \
  FIXED_BLOCK_DIM_BASE(64, ##__VA_ARGS__);   \
  FIXED_BLOCK_DIM_BASE(32, ##__VA_ARGS__)
#endif

struct SegmentOffsetIter {
  explicit SegmentOffsetIter(int num_cols) : num_cols_(num_cols) {}

  __host__ __device__ __forceinline__ int operator()(int idx) const {
#if !defined(__CUDA_ARCH__) && !defined(__HIP_DEVICE_COMPILE__)
    PADDLE_ENFORCE_LE_INT_MAX(static_cast<int64_t>(idx) * num_cols_,
                              "idx * num_cols_");
#endif
    return static_cast<int>(static_cast<int64_t>(idx) * num_cols_);
  }

  int num_cols_;
};

template <typename T>
struct Pair {
  __device__ __forceinline__ Pair() {}
  __device__ __forceinline__ Pair(T value, int id) : v(value), id(id) {}

  __device__ __forceinline__ void set(T value, int id) {
    this->v = value;
    this->id = id;
  }

  __device__ __forceinline__ void operator=(const Pair<T>& in) {
    v = in.v;
    id = in.id;
  }

  __device__ __forceinline__ bool operator<(const T value) const {
    return (static_cast<float>(v) < static_cast<float>(value));
  }

  __device__ __forceinline__ bool operator>(const T value) const {
    return (static_cast<float>(v) > static_cast<float>(value));
  }
  __device__ __forceinline__ bool operator<(const Pair<T>& in) const {
    return (static_cast<float>(v) < static_cast<float>(in.v)) ||
           ((static_cast<float>(v) == static_cast<float>(in.v)) &&
            (id > in.id));
  }

  __device__ __forceinline__ bool operator>(const Pair<T>& in) const {
    return (static_cast<float>(v) > static_cast<float>(in.v)) ||
           ((static_cast<float>(v) == static_cast<float>(in.v)) &&
            (id < in.id));
  }

  T v;
  int id;
};

int GetBlockSize(int vocab_size) {
  if (vocab_size > 512) {
    return 1024;
  } else if (vocab_size > 256) {
    return 512;
  } else if (vocab_size > 128) {
    return 256;
  } else if (vocab_size > 64) {
    return 128;
  } else {
    return 64;
  }
}

inline int64_t div_up(int64_t a, int64_t n) { return (a + n - 1) / n; }

template <typename T>
__device__ __forceinline__ void AddTo(Pair<T> topk[],
                                      const Pair<T>& p,
                                      int beam_size) {
  for (int k = beam_size - 2; k >= 0; k--) {
    if (topk[k] < p) {
      topk[k + 1] = topk[k];
    } else {
      topk[k + 1] = p;
      return;
    }
  }
  topk[0] = p;
}

template <typename T, int BlockSize>
__device__ __forceinline__ void GetTopK(
    Pair<T> topk[], const T* src, int idx, int dim, int beam_size) {
  while (idx < dim) {
    if (topk[beam_size - 1] < src[idx]) {
      Pair<T> tmp(src[idx], idx);
      AddTo<T>(topk, tmp, beam_size);
    }
    idx += BlockSize;
  }
}

template <typename T, int BlockSize>
__device__ __forceinline__ void GetTopK(Pair<T> topk[],
                                        const T* src,
                                        int idx,
                                        int dim,
                                        const Pair<T>& max,
                                        int beam_size) {
  while (idx < dim) {
    if (topk[beam_size - 1] < src[idx]) {
      Pair<T> tmp(src[idx], idx);
      if (tmp < max) {
        AddTo<T>(topk, tmp, beam_size);
      }
    }
    idx += BlockSize;
  }
}

template <typename T, int MaxLength, int BlockSize>
__device__ __forceinline__ void ThreadGetTopK(Pair<T> topk[],
                                              int* beam,
                                              int beam_size,
                                              const T* src,
                                              bool* firstStep,
                                              bool* is_empty,
                                              Pair<T>* max,
                                              int dim,
                                              const int tid) {
  if (*beam > 0) {
    int length = (*beam) < beam_size ? *beam : beam_size;
    if (*firstStep) {
      *firstStep = false;
      GetTopK<T, BlockSize>(topk, src, tid, dim, length);
    } else {
      for (int k = 0; k < MaxLength; k++) {
        if (k < MaxLength - (*beam)) {
          topk[k] = topk[k + *beam];
        } else {
          topk[k].set(std::numeric_limits<T>::min(), -1);
        }
      }
      if (!(*is_empty)) {
        GetTopK<T, BlockSize>(
            topk + MaxLength - *beam, src, tid, dim, *max, length);
      }
    }

    *max = topk[MaxLength - 1];
    if ((*max).id == -1) *is_empty = true;
    *beam = 0;
  }
}

template <typename T>
__forceinline__ __device__ Pair<T> WarpReduce(Pair<T> input) {
#pragma unroll
  for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
    T tmp_val = backends::gpu::CudaShuffleDownSync(
        FINAL_MASK, input.v, offset, WARP_SIZE);
    int tmp_id = backends::gpu::CudaShuffleDownSync(
        FINAL_MASK, input.id, offset, WARP_SIZE);
    if (static_cast<float>(input.v) < static_cast<float>(tmp_val)) {
      input.v = tmp_val;
      input.id = tmp_id;
    }
  }
  return input;
}

template <typename T, int MaxLength, int BlockSize>
__device__ __forceinline__ void BlockReduce(Pair<T> shared_max[],
                                            Pair<T> topk[],
                                            Pair<T> beam_max[],
                                            int* beam,
                                            int* k,
                                            int* count,
                                            const int tid,
                                            const int wid,
                                            const int lane) {
  while (true) {
    __syncthreads();
    Pair<T> input_now = topk[0];
    input_now = WarpReduce(input_now);

    if (lane == 0) {
      shared_max[wid] = input_now;
    }
    __syncthreads();
    input_now = (tid < BlockSize / WARP_SIZE)
                    ? shared_max[lane]
                    : Pair<T>(std::numeric_limits<T>::min(), -1);
    if (wid == 0) {
      input_now = WarpReduce(input_now);
      if (lane == 0) shared_max[0] = input_now;
    }
    __syncthreads();
    if (tid == 0) {
      beam_max[*count] = shared_max[0];
      (*count)++;
    }
    int tid_max = shared_max[0].id % BlockSize;
    if (tid == tid_max) {
      (*beam)++;
    }
    if (--(*k) == 0) break;
    __syncthreads();

    if (tid == tid_max) {
      if (*beam < MaxLength) {
        topk[0] = topk[*beam];
      }
    }

    if (MaxLength < 5) {
      if (*beam >= MaxLength) break;
    } else {
#ifdef PADDLE_WITH_HIP
      uint64_t mask = 0u;
      mask = __ballot(true);
      if (tid_max / WARP_SIZE == wid) {
        if (__shfl_down(*beam, tid_max % WARP_SIZE, WARP_SIZE) == MaxLength)
          break;
      }
#else
      unsigned mask = 0u;
      mask = __ballot_sync(FINAL_MASK, true);
      if (tid_max / WARP_SIZE == wid) {
        if (__shfl_down_sync(
                FINAL_MASK, *beam, tid_max % WARP_SIZE, WARP_SIZE) == MaxLength)
          break;
      }
#endif
    }
  }
}

template <typename T>
__device__ inline T exponential_transform(T val, T lambda) {
#if defined(__NVCC__) || defined(__HIPCC__)
  T log = -std::numeric_limits<T>::epsilon() / 2;
  if (val < static_cast<T>(1.) - std::numeric_limits<T>::epsilon() / 2) {
    if (std::is_same<T, double>::value) {
      log = logf(val);
    } else {
      log = __logf(val);
    }
  }
  return static_cast<T>(-1.0) / lambda * log;
#else
  return static_cast<T>(-1.0) / lambda * std::log(static_cast<T>(1.0) - val);
#endif
}

template <typename T, int MaxLength, int TopPBeamTopK, int BlockSize>
__global__ void KeMatrixTopPBeamTopK(const T* src,
                                     const T* threshold,
                                     GPU(randState_t) * states,
                                     T* top_ps,
                                     int64_t* out_id,  // topk id
                                     T* out_val,       // topk val
                                     int64_t* topk_ids,
                                     T* topk_scores,
                                     int vocab_size,
                                     int* count_iter,
                                     int* count_iter_begin,
                                     const int k,
                                     const bool need_batch_random) {
  const int tid = threadIdx.x;
  const int wid = tid / WARP_SIZE;
  const int lane = tid % WARP_SIZE;
  const int bid = blockIdx.x;
  const float threshold_now =
      threshold ? static_cast<float>(threshold[bid]) : 0.f;

  int top_num = TopPBeamTopK;
  float top_p_num = static_cast<float>(top_ps[bid]);
  const int64_t offset = static_cast<int64_t>(bid) * vocab_size;
  int64_t* topk_ids_now = nullptr;
  T* topk_scores_now = nullptr;
  if (k > 0) {
    topk_ids_now = topk_ids + static_cast<int64_t>(bid) * k;
    topk_scores_now = topk_scores + static_cast<int64_t>(bid) * k;
  }

  __shared__ Pair<T> shared_max[BlockSize / WARP_SIZE];
  __shared__ Pair<T> beam_max[TopPBeamTopK];

  Pair<T> topk[MaxLength];
  int beam = MaxLength;
  Pair<T> max;
  bool is_empty = false;
  bool firststep = true;
  __shared__ int count;

  if (tid == 0) {
    count = 0;
  }

  for (int j = 0; j < MaxLength; j++) {
    topk[j].set(std::numeric_limits<T>::min(), -1);
  }

  while (top_num) {
    ThreadGetTopK<T, MaxLength, BlockSize>(topk,
                                           &beam,
                                           TopPBeamTopK,
                                           src + offset,
                                           &firststep,
                                           &is_empty,
                                           &max,
                                           vocab_size,
                                           tid);
    BlockReduce<T, MaxLength, BlockSize>(
        shared_max, topk, beam_max, &beam, &top_num, &count, tid, wid, lane);
  }
  if (tid == 0) {
    count_iter_begin[bid] = count_iter[bid];
    float top_p = top_ps[bid];
    float sum_prob = 0.0f;
    bool flag = false;
    float max_val = 0.f;
    int max_id = -1;
    for (int i = 0; i < TopPBeamTopK; i++) {
      if (i < k) {
        topk_ids_now[i] = static_cast<int64_t>(beam_max[i].id);
        topk_scores_now[i] = beam_max[i].v;
      }
      if (!flag) {
        float val = static_cast<float>(beam_max[i].v);
        sum_prob += val;
        float random_ratio =
            exponential_transform(GPU(rand_uniform)(states + bid), 1.0f);
        float random_val = (val >= threshold_now ? val : 0.f) / random_ratio;
        if (max_val < random_val) {
          max_val = random_val;
          max_id = i;
        }
        if (sum_prob >= top_p) {
          flag = true;
          count_iter_begin[bid] += 1;
          if (max_id == -1) {
            // don't sample low score token
            out_id[bid] = static_cast<int64_t>(beam_max[0].id);
            out_val[bid] = beam_max[0].v;
          } else {
            out_id[bid] = static_cast<int64_t>(beam_max[max_id].id);
            out_val[bid] = beam_max[max_id].v;
          }
        }
      }
      if (flag && i >= k - 1) {
        break;
      }
    }
  }
}

template <typename T, int MaxLength, int TopPBeamTopK, int BlockSize>
__global__ void KeMatrixTopPBeamTopKFt(const T* src,
                                       const T* threshold,
                                       GPU(randState_t) * states,
                                       T* top_ps,
                                       int64_t* out_id,  // topk id
                                       T* out_val,       // topk val
                                       int64_t* topk_ids,
                                       T* topk_scores,
                                       int vocab_size,
                                       int* count_iter,
                                       int* count_iter_begin,
                                       const int k,
                                       const bool need_batch_random) {
  const int tid = threadIdx.x;
  const int wid = tid / WARP_SIZE;
  const int lane = tid % WARP_SIZE;
  const int bid = blockIdx.x;
  const float threshold_now =
      threshold ? static_cast<float>(threshold[bid]) : 0.f;

  int top_num = TopPBeamTopK;
  float top_p_num = static_cast<float>(top_ps[bid]);
  int64_t* topk_ids_now = nullptr;
  T* topk_scores_now = nullptr;
  if (k > 0) {
    topk_ids_now = topk_ids + bid * k;
    topk_scores_now = topk_scores + bid * k;
  }

  __shared__ Pair<T> shared_max[BlockSize / WARP_SIZE];
  __shared__ Pair<T> beam_max[TopPBeamTopK];

  Pair<T> topk[MaxLength];
  int beam = MaxLength;
  Pair<T> max;
  bool is_empty = false;
  bool firststep = true;
  __shared__ int count;

  if (tid == 0) {
    count = 0;
  }

  for (int j = 0; j < MaxLength; j++) {
    topk[j].set(std::numeric_limits<T>::min(), -1);
  }

  while (top_num) {
    ThreadGetTopK<T, MaxLength, BlockSize>(
        topk,
        &beam,
        TopPBeamTopK,
        src + static_cast<int64_t>(bid) * vocab_size,
        &firststep,
        &is_empty,
        &max,
        vocab_size,
        tid);
    BlockReduce<T, MaxLength, BlockSize>(
        shared_max, topk, beam_max, &beam, &top_num, &count, tid, wid, lane);
  }
  if (tid == 0) {
    count_iter_begin[bid] = count_iter[bid];
    float rand_top_p = GPU(rand_uniform)(states + bid) * top_p_num;
    top_ps[bid] = (T)rand_top_p;
    float sum_prob = 0.0f;
    bool flag = false;
    for (int i = 0; i < TopPBeamTopK; i++) {
      if (i < k) {
        topk_ids_now[i] = static_cast<int64_t>(beam_max[i].id);
        topk_scores_now[i] = beam_max[i].v;
      }
      if (!flag) {
        float val = static_cast<float>(beam_max[i].v);
        sum_prob += val;
#ifdef DEBUG_TOPP
        printf("bi: %d, top_p: %f, rand_top_p: %f, sum_prob: %f\n",
               bid,
               top_p_num,
               rand_top_p,
               sum_prob);
#endif
        if (sum_prob >= rand_top_p) {
          flag = true;
          count_iter_begin[bid] += 1;
          if (val < threshold_now) {
            // don't sample low score token
            int start_id = i == 0 ? 0 : i - 1;
            for (int j = start_id; j >= 0; j--) {
              float val_now = static_cast<float>(beam_max[j].v);
              if (val_now >= threshold_now || j == 0) {
                out_id[bid] = static_cast<int64_t>(beam_max[j].id);
                out_val[bid] = beam_max[j].v;
                break;
              }
            }
          } else {
            out_id[bid] = static_cast<int64_t>(beam_max[i].id);
            out_val[bid] = beam_max[i].v;
          }
        }
      }
      if (flag && i >= k - 1) {
        break;
      }
    }
  }
}

__global__ void SetCountIter(int* count_iter, int num) {
  int tid = threadIdx.x;
  int bid = blockIdx.x;
  int idx = bid * blockDim.x + tid;
  for (int64_t i = idx; i < num; i += static_cast<int64_t>(gridDim.x) *
                                      static_cast<int64_t>(blockDim.x)) {
    count_iter[i] = i;
  }
}

template <typename T>
__global__ void FillIndex(T* indices, T num_rows, T num_cols) {
  int col_id = threadIdx.x;
  int row_id = blockIdx.x;

  for (T j = row_id; j < num_rows; j += gridDim.x) {
    for (T i = col_id; i < num_cols; i += blockDim.x) {
      indices[j * num_cols + i] = i;
    }
  }
}

template <typename T, typename Context, int TopKMaxLength, int TopPBeamTopK>
void DispatchKeMatrixTopPBeamTopK(const Context& dev_ctx,
                                  const T* src,
                                  const T* threshold,
                                  GPU(randState_t) * states,
                                  T* top_ps,
                                  int64_t* out_id,  // topk id
                                  T* out_val,       // topk val
                                  int64_t* topk_ids,
                                  T* topk_scores,
                                  int vocab_size,
                                  int* count_iter,
                                  int* count_iter_begin,
                                  const int k,
                                  const int bs,
                                  const bool need_batch_random,
                                  const std::string& mode) {
  int BlockSize = GetBlockSize(vocab_size);
  if (mode == "truncated") {
    switch (BlockSize) {
      FIXED_BLOCK_DIM(
          KeMatrixTopPBeamTopKFt<T, TopKMaxLength, TopPBeamTopK, kBlockDim>
          <<<bs, kBlockDim, 0, dev_ctx.stream()>>>(src,
                                                   threshold,
                                                   states,
                                                   top_ps,
                                                   out_id,
                                                   out_val,
                                                   topk_ids,
                                                   topk_scores,
                                                   vocab_size,
                                                   count_iter,
                                                   count_iter_begin,
                                                   k,
                                                   need_batch_random));
      default:
        PD_THROW(
            "the input data shape has error in the topp_beam_topk kernel.");
    }
  } else {
    switch (BlockSize) {
      FIXED_BLOCK_DIM(
          KeMatrixTopPBeamTopK<T, TopKMaxLength, TopPBeamTopK, kBlockDim>
          <<<bs, kBlockDim, 0, dev_ctx.stream()>>>(src,
                                                   threshold,
                                                   states,
                                                   top_ps,
                                                   out_id,
                                                   out_val,
                                                   topk_ids,
                                                   topk_scores,
                                                   vocab_size,
                                                   count_iter,
                                                   count_iter_begin,
                                                   k,
                                                   need_batch_random));
      default:
        PD_THROW(
            "the input data shape has error in the topp_beam_topk kernel.");
    }
  }
}

struct BlockPrefixCallbackOp {
  // Running prefix
  float running_total;
  // Constructor
  __device__ BlockPrefixCallbackOp(float running_total)
      : running_total(running_total) {}
  // Callback operator to be entered by the first warp of threads in the block.
  // Thread-0 is responsible for returning a value for seeding the block-wide
  // scan.
  __device__ float operator()(float block_aggregate) {
    float old_prefix = running_total;
    running_total += block_aggregate;
    return old_prefix;
  }
};

template <typename T>
__device__ T max_func(const T a, const T b) {
  return a > b ? a : b;
}

template <typename T>
struct MaxOp {
  __device__ __forceinline__ T operator()(const T& a, const T& b) const {
    return max_func(a, b);
  }
};

template <typename T, int BLOCK_SIZE>
__global__ void topp_sampling(T* sorted_probs,
                              int64_t* sorted_id,
                              T* out_val,
                              int64_t* out_id,
                              const T* top_ps,
                              const T* threshold,
                              const int64_t* infer_seed,
                              GPU(randState_t) * states,
                              const int p_num,
                              const uint64_t seed,
                              const int vocab_size,
                              const bool need_batch_random,
                              int* count_iter,
                              int* count_iter_begin) {
  __shared__ int stop_shared;
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;
  const float p_t = static_cast<float>(top_ps[bid]);
  const float threshold_now =
      threshold ? static_cast<float>(threshold[bid]) : 0.f;
  uint64_t seed_now;
  GPU(randState_t) rand_state;
  const int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (infer_seed) {
    seed_now = static_cast<uint64_t>(infer_seed[bid]);
    GPU(rand_init)(seed_now, tid, 0, &rand_state);
  } else {
    seed_now = seed;
    GPU(rand_init)(seed_now, global_idx, 0, &rand_state);
  }
  if (tid == 0) {
    stop_shared = 0;
  }
  if (count_iter_begin[bid] == count_iter[bid + 1]) {
    // topk
    return;
  }

  typedef cub::BlockScan<float, BLOCK_SIZE> BlockScan;
  typedef cub::BlockReduce<Pair<T>, BLOCK_SIZE> BlockReduce;
  __shared__ typename BlockScan::TempStorage temp_storage;
  __shared__ typename BlockReduce::TempStorage temp_storage_reduce;

  // Initialize running total
  BlockPrefixCallbackOp prefix_op(0);

  int64_t offset = static_cast<int64_t>(bid) * vocab_size;
#ifdef DEBUG_TOPP
  if (tid == 0) {
    printf(
        "first_elem1_1: %f, first_elem1_2: %f, first_id1_1: %d, first_id1_2: "
        "%d\n",
        static_cast<float>(sorted_probs[offset]),
        static_cast<float>(sorted_probs[offset + 1]),
        static_cast<int>(sorted_id[offset]),
        static_cast<int>(sorted_id[offset + 1]));
  }
#endif
  int end = ((vocab_size + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
  int i_activate = 0;
  float thread_offset = 0;
  Pair<T> max_thread_pair(static_cast<T>(0.), -1);
  for (int i = tid; i < end; i += BLOCK_SIZE) {
    float thread_count =
        (i < vocab_size) ? static_cast<float>(sorted_probs[offset + i]) : 0.f;
    BlockScan(temp_storage)
        .InclusiveSum(thread_count, thread_offset, prefix_op);

    if (thread_offset < p_t ||
        (thread_offset >= p_t && thread_offset - thread_count < p_t)) {
      float random_ratio =
          exponential_transform(GPU(rand_uniform)(&rand_state), 1.0f);
      float tmp_val =
          (thread_count >= threshold_now ? thread_count : 0.f) / random_ratio;
      if (static_cast<float>(max_thread_pair.v) < tmp_val) {
        max_thread_pair.set(static_cast<T>(tmp_val), i);
      }
#ifdef DEBUG_TOPP
      if (i < 10) {
        printf(
            "tid: %d, i: %d, random_ratio: %f, thread_count: %f, tmp_val: %f, "
            "max_thread_pair.v: %f, max_thread_pair.id: %d\n",
            tid,
            i,
            random_ratio,
            thread_count,
            tmp_val,
            max_thread_pair.v,
            static_cast<int>(max_thread_pair.id));
      }
#endif
    }
#ifdef DEBUG_TOPP
    printf("tid: %d, thread_count: %f, thread_offset: %f\n",
           tid,
           thread_count,
           thread_offset);
#endif
#ifdef PADDLE_WITH_HIP
    uint64_t activate_mask = __ballot(p_t <= thread_offset);
#else
    uint32_t activate_mask = __ballot_sync(FINAL_MASK, p_t <= thread_offset);
#endif

    i_activate = i;
    if (activate_mask != 0) {
      if (lane_id == 0) {
        atomicAdd(&stop_shared, 1);
      }
    }
    __syncthreads();
    if (stop_shared > 0) {
      break;
    }
  }
  __syncthreads();

  Pair<T> max_pair = BlockReduce(temp_storage_reduce)
                         .Reduce(max_thread_pair, MaxOp<Pair<T>>());
  if (tid == 0) {
    if (max_pair.id == -1) {
      max_pair.id = 0;
    }
#ifdef DEBUG_TOPP
    printf("max_id: %d, max_val: %f\n",
           static_cast<int>(max_pair.id),
           static_cast<float>(max_pair.v));
#endif
    out_id[bid] = sorted_id[offset + max_pair.id];
    out_val[bid] = sorted_probs[offset + max_pair.id];
  }
}

template <typename T, int BLOCK_SIZE>
__global__ void topp_sampling_ft(T* sorted_probs,
                                 int64_t* sorted_id,
                                 T* out_val,
                                 int64_t* out_id,
                                 const T* top_ps,
                                 const T* threshold,
                                 GPU(randState_t) * states,
                                 const int p_num,
                                 const int vocab_size,
                                 const bool need_batch_random,
                                 int* count_iter,
                                 int* count_iter_begin) {
  __shared__ int stop_shared;
  __shared__ float rand_p;
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;
  const float p_t = static_cast<float>(top_ps[bid]);
  const float threshold_now =
      threshold ? static_cast<float>(threshold[bid]) : 0.f;
  if (tid == 0) {
    stop_shared = 0;
    rand_p = p_t;
#ifdef DEBUG_TOPP
    printf("bi: %d, p: %f\n", bid, rand_p);
#endif
  }
  if (count_iter_begin[bid] == count_iter[bid + 1]) {
    // topk
    return;
  }

  typedef cub::BlockScan<float, BLOCK_SIZE> BlockScan;
  typedef cub::BlockReduce<int, BLOCK_SIZE> BlockReduce;
  __shared__ typename BlockScan::TempStorage temp_storage;
  __shared__ typename BlockReduce::TempStorage temp_storage_reduce;
#ifdef PADDLE_WITH_HIP
  __shared__ uint64_t selected_shared[NUM_WARPS];
#else
  __shared__ uint32_t selected_shared[NUM_WARPS];
#endif
  int threshold_id = 0;

  // Initialize running total
  BlockPrefixCallbackOp prefix_op(0);

  if (lane_id == 0) {
    selected_shared[warp_id] = 0;
  }
  __syncthreads();

  int64_t offset = static_cast<int64_t>(bid) * vocab_size;
#ifdef DEBUG_TOPP
  if (tid == 0) {
    printf(
        "first_elem1_1: %f, first_elem1_2: %f, first_id1_1: %d, first_id1_2: "
        "%d\n",
        static_cast<float>(sorted_probs[offset]),
        static_cast<float>(sorted_probs[offset + 1]),
        static_cast<int>(sorted_id[offset]),
        static_cast<int>(sorted_id[offset + 1]));
  }
#endif
  int end = ((vocab_size + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
  int i_activate = 0;
  float thread_offset = 0;
  for (int i = tid; i < end; i += BLOCK_SIZE) {
    float thread_count =
        (i < vocab_size) ? static_cast<float>(sorted_probs[offset + i]) : 0.f;
    if (i < vocab_size && thread_count >= threshold_now) {
      threshold_id = i;
    }
    BlockScan(temp_storage)
        .InclusiveSum(thread_count, thread_offset, prefix_op);

#ifdef PADDLE_WITH_HIP
    uint64_t activate_mask = __ballot(rand_p <= thread_offset);
#else
    uint32_t activate_mask = __ballot_sync(FINAL_MASK, rand_p <= thread_offset);
#endif

    i_activate = i;
    if (activate_mask != 0) {
      if (lane_id == 0) {
        atomicAdd(&stop_shared, 1);
        selected_shared[warp_id] = activate_mask;
      }
    }
    __syncthreads();
    if (stop_shared > 0) {
      break;
    }
  }
  __syncthreads();
  if (stop_shared == 0) {
    if (tid == 0) {
      out_id[bid] = sorted_id[offset];
      out_val[bid] = sorted_probs[offset];
#ifdef DEBUG_TOPP
      printf("stop_shared: %d, out_id: %d, out_val: %f\n",
             static_cast<int>(stop_shared),
             static_cast<int>(out_id[bid]),
             static_cast<float>(out_val[bid]));
#endif
    }
    return;
  }
#ifdef DEBUG_TOPP
  if (tid == 0) {
    printf(
        "first_elem2_1: %f, first_elem2_2: %f, first_id2_1: %d, first_id2_2: "
        "%d\n",
        static_cast<float>(sorted_probs[offset]),
        static_cast<float>(sorted_probs[offset + 1]),
        static_cast<int>(sorted_id[offset]),
        static_cast<int>(sorted_id[offset + 1]));
  }
#endif
  bool skip = (selected_shared[warp_id] > 0) ? false : true;
  for (int i = 0; i < warp_id; i++) {
    if (selected_shared[i] != 0) {
      // If the previous has stopped, skip the current warp
      skip = true;
    }
  }
  if (!skip) {
#ifdef PADDLE_WITH_HIP
    int active_lane_id =
        WARP_SIZE - __popcll(selected_shared[warp_id]);  // first not 0
#else
    int active_lane_id =
        WARP_SIZE - __popc(selected_shared[warp_id]);  // first not 0
#endif
    if (lane_id == active_lane_id) {
      float val = static_cast<float>(sorted_probs[offset + i_activate]);
#ifdef DEBUG_TOPP
      printf(
          "active_lane_id: %d, i_activate: %d.\n", active_lane_id, i_activate);
      for (int i = 0; i < active_lane_id; i++) {
        printf("p %d, value: %f\n",
               i,
               static_cast<float>(sorted_probs[offset + i]));
      }
#endif
      if (val < threshold_now) {
        // don't sample low score token
        int max_id =
            BlockReduce(temp_storage_reduce).Reduce(threshold_id, MaxOp<int>());
#ifdef PADDLE_WITH_HIP
        hiprandStatePhilox4_32_10_t rng;
        hiprand_init(bid * blockDim.x + tid, tid, 0, &rng);
        int random_id = hiprand(&rng) % (max_id + 1);
#else
        curandStatePhilox4_32_10_t rng;
        curand_init(bid * blockDim.x + tid, tid, 0, &rng);
        int random_id = curand(&rng) % (max_id + 1);
#endif
        out_id[bid] = sorted_id[offset + random_id];
        out_val[bid] = sorted_probs[offset + random_id];
      } else {
        out_id[bid] = sorted_id[offset + i_activate];
        out_val[bid] = sorted_probs[offset + i_activate];
      }
    }
  }
}

template <typename T, typename Context>
void DispatchTopPSampling(const Context& dev_ctx,
                          T* sorted_probs,
                          int64_t* sorted_id,
                          T* out_val,
                          int64_t* out_id,
                          const T* top_ps,
                          const T* threshold,
                          const int64_t* infer_seed,
                          GPU(randState_t) * states,
                          const int p_num,
                          const int vocab_size,
                          const int bs,
                          const uint64_t seed,
                          const bool need_batch_random,
                          int* count_iter,
                          int* count_iter_begin,
                          const std::string& mode) {
  int BlockSize = GetBlockSize(vocab_size);
  if (mode == "truncated") {
    switch (BlockSize) {
      FIXED_BLOCK_DIM(
          topp_sampling_ft<T, kBlockDim>
          <<<bs, kBlockDim, 0, dev_ctx.stream()>>>(sorted_probs,
                                                   sorted_id,
                                                   out_val,
                                                   out_id,
                                                   top_ps,
                                                   threshold,
                                                   states,
                                                   p_num,
                                                   vocab_size,
                                                   need_batch_random,
                                                   count_iter,
                                                   count_iter_begin));
      default:
        PD_THROW("the input data shape has error in the topp_sampling kernel.");
    }
  } else {
    switch (BlockSize) {
      FIXED_BLOCK_DIM(
          topp_sampling<T, kBlockDim>
          <<<bs, kBlockDim, 0, dev_ctx.stream()>>>(sorted_probs,
                                                   sorted_id,
                                                   out_val,
                                                   out_id,
                                                   top_ps,
                                                   threshold,
                                                   infer_seed,
                                                   states,
                                                   p_num,
                                                   seed,
                                                   vocab_size,
                                                   need_batch_random,
                                                   count_iter,
                                                   count_iter_begin));
      default:
        PD_THROW("the input data shape has error in the topp_sampling kernel.");
    }
  }
}

__global__ void setup_kernel(GPU(randState_t) * state,
                             int64_t* seed,
                             const int bs) {
  int64_t idx =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  for (int64_t i = idx; i < bs; i += static_cast<int64_t>(gridDim.x) *
                                     static_cast<int64_t>(blockDim.x)) {
    GPU(rand_init)(static_cast<uint64_t>(seed[i]), 0, 0, &state[i]);
  }
}

__global__ void setup_kernel(GPU(randState_t) * state,
                             const uint64_t seed,
                             const uint64_t offset,
                             const int bs,
                             const bool need_batch_random) {
  int64_t idx =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);
  for (int64_t i = idx; i < bs; i += static_cast<int64_t>(gridDim.x) *
                                     static_cast<int64_t>(blockDim.x)) {
    if (need_batch_random) {
      GPU(rand_init)(seed, i, offset, &state[i]);
    } else {
      GPU(rand_init)(seed, 0, offset, &state[i]);
    }
  }
}

template <typename T>
T* SafeGetTensorPtr(const DenseTensor& t) {
  return const_cast<T*>(t.data<T>());
}

template <typename T>
T* SafeGetTensorPtr(const DenseTensor* t) {
  return t ? SafeGetTensorPtr<T>(*t) : nullptr;
}

template <typename T>
T* SafeGetTensorPtr(const optional<DenseTensor>& t) {
  return t ? SafeGetTensorPtr<T>(t.get()) : nullptr;
}

template <typename T, typename Context>
void TopPSamplingKernel(const Context& dev_ctx,
                        const DenseTensor& x,
                        const DenseTensor& ps,
                        const optional<DenseTensor>& threshold,
                        const optional<DenseTensor>& topp_seed,
                        int64_t seed,
                        int k,
                        const std::string& mode,
                        DenseTensor* out,
                        DenseTensor* ids,
                        DenseTensor* topk_scores,
                        DenseTensor* topk_ids) {
  typedef DataTypeTraits<T> traits_;
  typedef typename traits_::DataType DataType_;
  auto cu_stream = dev_ctx.stream();
  const auto* input = &x;
  // get the input dims
  const auto& in_dims = input->dims();
  int64_t p_num = ps.numel();
  int64_t bs = in_dims[0];
  // TODO(large-tensor): downstream functors may still use int
  PADDLE_ENFORCE_LE_INT_MAX(p_num, "p_num");
  PADDLE_ENFORCE_LE_INT_MAX(bs + 1, "bs + 1");
  PADDLE_ENFORCE_LE_INT_MAX(in_dims[1], "vocab_size");
  int64_t num_items64 = bs * in_dims[1];
  PADDLE_ENFORCE_LE_INT_MAX(num_items64, "bs * vocab_size");
  const int p_num_int = static_cast<int>(p_num);
  const int vocab_size = static_cast<int>(in_dims[1]);
  const int num_segments = static_cast<int>(bs);
  const int num_segments_with_end = static_cast<int>(bs + 1);
  const int num_items = static_cast<int>(num_items64);
  T* out_ptr = dev_ctx.template Alloc<T>(out);
  int64_t* ids_ptr = dev_ctx.template Alloc<int64_t>(ids);
  T* topk_scores_data = nullptr;
  int64_t* topk_ids_data = nullptr;
  if (k > 0) {
    topk_scores_data = dev_ctx.template Alloc<T>(topk_scores);
    topk_ids_data = dev_ctx.template Alloc<int64_t>(topk_ids);
  }

  DenseTensor ps_now;
  ps_now.Resize({bs, 1});
  dev_ctx.template Alloc<T>(&ps_now);
  Copy(dev_ctx, ps, dev_ctx.GetPlace(), false, &ps_now);

  DenseTensor inds_input;
  inds_input.Resize({bs, vocab_size});
  dev_ctx.template Alloc<int64_t>(&inds_input);

  DenseTensor sorted_out;
  sorted_out.Resize({bs, vocab_size});
  dev_ctx.template Alloc<T>(&sorted_out);

  DenseTensor sorted_id;
  sorted_id.Resize({bs, vocab_size});
  dev_ctx.template Alloc<int64_t>(&sorted_id);

  int BlockSize = GetBlockSize(vocab_size);

  switch (BlockSize) {
    FIXED_BLOCK_DIM(FillIndex<int64_t>
                    <<<num_segments, kBlockDim, 0, cu_stream>>>(
                        inds_input.data<int64_t>(), bs, vocab_size));
    default:
      PD_THROW("the input data shape has error in the FillIndex kernel.");
  }
  int64_t* infer_seed = SafeGetTensorPtr<int64_t>(topp_seed);

  GPU(randState_t) * states{nullptr};
  phi::Allocator::AllocationPtr rand_states_buf{nullptr};
  rand_states_buf = phi::memory_utils::Alloc(
      dev_ctx.GetPlace(),
      bs * sizeof(GPU(randState_t)),
      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
  states = reinterpret_cast<GPU(randState_t)*>(rand_states_buf->ptr());

  uint64_t seed_now = seed;
  uint64_t offset = 0;
  bool need_batch_random = false;

  if (infer_seed) {
    setup_kernel<<<1, 256, 0, cu_stream>>>(states, infer_seed, num_segments);
  } else {
    if (seed_now == -1) {
      need_batch_random = true;
      auto gen_cuda = dev_ctx.GetGenerator();
      uint64_t increment = bs * BlockSize;
      auto seed_offset = gen_cuda->IncrementOffset(increment);
      seed_now = seed_offset.first;
      offset = seed_offset.second;
      setup_kernel<<<1, 256, 0, cu_stream>>>(
          states, seed_now, offset, num_segments, need_batch_random);
    } else {
      setup_kernel<<<1, 256, 0, cu_stream>>>(
          states, seed_now, offset, num_segments, need_batch_random);
    }
  }

  DenseTensor count_iter;
  count_iter.Resize({bs + 1});
  dev_ctx.template Alloc<int>(&count_iter);
  DenseTensor count_iter_begin;
  count_iter_begin.Resize({bs});
  dev_ctx.template Alloc<int>(&count_iter_begin);
  SetCountIter<<<1, 256, 0, cu_stream>>>(count_iter.data<int>(),
                                         num_segments_with_end);

  T* threshold_data = SafeGetTensorPtr<T>(threshold);

  constexpr int TopKMaxLength = 2;
  constexpr int TopPBeamTopK = 20;

  DispatchKeMatrixTopPBeamTopK<T, Context, TopKMaxLength, TopPBeamTopK>(
      dev_ctx,
      x.data<T>(),
      threshold_data,
      states,
      ps_now.data<T>(),
      ids_ptr,
      out_ptr,
      topk_ids_data,
      topk_scores_data,
      vocab_size,
      count_iter.data<int>(),
      count_iter_begin.data<int>(),
      k,
      num_segments,
      need_batch_random,
      mode);

  size_t temp_storage_bytes = 0;

  cub::TransformInputIterator<int, SegmentOffsetIter, int*>
      segment_offsets_t_begin(count_iter_begin.data<int>(),
                              SegmentOffsetIter(vocab_size));

  cub::TransformInputIterator<int, SegmentOffsetIter, int*>
      segment_offsets_t_end(count_iter.data<int>(),
                            SegmentOffsetIter(vocab_size));

  cub::DeviceSegmentedRadixSort::SortPairsDescending(
      nullptr,
      temp_storage_bytes,
      reinterpret_cast<DataType_*>(const_cast<T*>(x.data<T>())),
      reinterpret_cast<DataType_*>(const_cast<T*>(sorted_out.data<T>())),
      inds_input.data<int64_t>(),
      sorted_id.data<int64_t>(),
      num_items,
      num_segments,
      segment_offsets_t_begin,
      segment_offsets_t_end + 1,
      0,
      sizeof(T) * 8,
      cu_stream);

  temp_storage_bytes = div_up(temp_storage_bytes, 256) * 256;
  int64_t temp_size = temp_storage_bytes;
  DenseTensor temp_storage;
  temp_storage.Resize({temp_size});
  dev_ctx.template Alloc<uint8_t>(&temp_storage);

  cub::DeviceSegmentedRadixSort::SortPairsDescending(
      temp_storage.data<uint8_t>(),
      temp_storage_bytes,
      reinterpret_cast<DataType_*>(const_cast<T*>(x.data<T>())),
      reinterpret_cast<DataType_*>(const_cast<T*>(sorted_out.data<T>())),
      inds_input.data<int64_t>(),
      sorted_id.data<int64_t>(),
      num_items,
      num_segments,
      segment_offsets_t_begin,
      segment_offsets_t_end + 1,
      0,
      sizeof(T) * 8,
      cu_stream);

  DispatchTopPSampling<T>(dev_ctx,
                          sorted_out.data<T>(),
                          sorted_id.data<int64_t>(),
                          out_ptr,
                          ids_ptr,
                          ps_now.data<T>(),
                          threshold_data,
                          infer_seed,
                          states,
                          p_num_int,
                          vocab_size,
                          num_segments,
                          seed_now + offset,
                          need_batch_random,
                          count_iter.data<int>(),
                          count_iter_begin.data<int>(),
                          mode);
}
}  // namespace phi

#ifdef CUDA_BFLOAT16_AVAILABLE
PD_REGISTER_KERNEL(top_p_sampling,
                   GPU,
                   ALL_LAYOUT,
                   phi::TopPSamplingKernel,
                   float,
                   double,
                   int,
                   int64_t,
                   phi::float16,
                   phi::bfloat16) {}
#else
PD_REGISTER_KERNEL(top_p_sampling,
                   GPU,
                   ALL_LAYOUT,
                   phi::TopPSamplingKernel,
                   float,
                   double,
                   int,
                   int64_t,
                   phi::float16) {}
#endif