   }
 }
 
-static __device__ __forceinline__ platform::float16 exp_on_device(
-    platform::float16 x) {
-  return ::Eigen::numext::exp(x);
-}
-static __device__ __forceinline__ float exp_on_device(float x) {
-  return expf(x);
-}
-static __device__ __forceinline__ double exp_on_device(double x) {
-  return exp(x);
-}
-static __device__ __forceinline__ platform::float16 log_on_device(
-    platform::float16 x) {
-  return math::TolerableValue<platform::float16>()(::Eigen::numext::log(x));
-}
-static __device__ __forceinline__ float log_on_device(float x) {
-  return math::TolerableValue<float>()(logf(x));
-}
-static __device__ __forceinline__ double log_on_device(double x) {
-  return math::TolerableValue<double>()(log(x));
-}
-
-/** In the following codes, 3 CUDA kernels are implemented to calculate softmax
- * and loss **/
-/*
-  Supposing the x is `logits` and y is `labels`, the equations are as
-followings:
-  cross\_entropy_i = \sum_{j}[- y_i_j * log({e^{x_i_j}/\sum_{j}e^{x_i_j}})]
-        = \sum_{j}[- y_i_j * log({e^{x_i_j - max_i}/\sum_{j}e^{x_i_j-max_i}})]
-        = \sum_{j}[-y_i_j * (x_i_j - max_i - log\sum_{j}e^{x_i_j - max_i})]
-        = \sum_{j}[-y_i_j * (x_i_j - max_i - logDiffMaxSum_i)]
-        = \sum_{j}(-y_i_j * tmp_i_j)
-  softmax_i_j = e^{tmp_i_j}
-where:
-  max_i = \max_{j}{x_i_j}
-  logDiffMaxSum_i = log\sum_{j}e^{x_i_j - max_i}
-  tmp_i_j = x_i_j - max_i - logDiffMaxSum_i
-Therefore, the calculation can be separated into 3 steps:
-Step 1: row-wise operation to calculate max_i
-Step 2: row-wise operation to calculate logDiffMaxSum_i
-Step 3: calculate tmp_i_j, and finally get softmax_i_j and cross\_entropy_i
-To save memory, we can share memory among max_i, logDiffMaxSum_i and
-cross\_entropy_i.
-In this way, the 3 steps should be changed to:
-Step 1 (RowReductionForMax): row-wise operation to calculate max_i
-Step 2 (RowReductionForDiffMaxSum): calculate immediate result of softmax'_i_j =
-x_i_j - max_i, and row-wise operation to calculate logDiffMaxSum_i
-Step 3 (RowReductionForSoftmaxAndCrossEntropy): calculate tmp_i_j = softmax'_i_j
-- logDiffMaxSum_i, and finally get softmax_i_j and cross\_entropy_i
-*/
-
-// There are 3 kinds of reduce algorithms in cub:
-// BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY
-// BLOCK_REDUCE_RAKING
-// BLOCK_REDUCE_WARP_REDUCTIONS (default)
-template <typename T, int BlockDim>
-using BlockReduce =
-    cub::BlockReduce<T, BlockDim /*, cub::BLOCK_REDUCE_WARP_REDUCTIONS*/>;
-
-template <typename T, int BlockDim>
-using BlockReduceTempStorage = typename BlockReduce<T, BlockDim>::TempStorage;
-
-// Make sure that BlockDim <= axis_dim
-// This kernel is used to calculate the max element of each row
-template <typename T, int BlockDim>
-static __global__ void RowReductionForMax(const T* logits_data, T* max_data,
-                                          int64_t d, int axis_dim) {
-  __shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
-
-  // logits_data view as [n, axis_dim, remain]
-  // max_data view as [n, 1, remain]
-  // blockDim = n * remain, split blockIdx to idx_n and idx_remain
-  int64_t remain = d / axis_dim;
-  int64_t idx_n = blockIdx.x / remain;
-  int64_t idx_remain = blockIdx.x % remain;
-  int64_t beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int64_t end_idx = (idx_n + 1) * d;
-
-  int64_t step = BlockDim * remain;
-  T cur_max = logits_data[beg_idx];
-  beg_idx += step;
-  while (beg_idx < end_idx) {
-    if (cur_max < logits_data[beg_idx]) {
-      cur_max = logits_data[beg_idx];
-    }
-    beg_idx += step;
-  }
-
-  cur_max = BlockReduce<T, BlockDim>(temp_storage).Reduce(cur_max, cub::Max());
-
-  if (threadIdx.x == 0) max_data[blockIdx.x] = cur_max;
-}
-
-// Make sure that BlockDim <= axis_dim
-template <typename T, int BlockDim, bool CalculateLogSoftmax = false>
-static __global__ void RowReductionForDiffMaxSum(const T* logits_data,
-                                                 T* max_data, T* softmax,
-                                                 int64_t d, int axis_dim) {
-  __shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
-
-  // logits, softmax data view as [n, axis_dim, remain]
-  // max_data view as [n, 1, remain]
-  // blockDim = n * remain, split blockIdx to idx_n and idx_remain
-  int64_t remain = d / axis_dim;
-  int64_t idx_n = blockIdx.x / remain;
-  int64_t idx_remain = blockIdx.x % remain;
-  int64_t beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int64_t end_idx = (idx_n + 1) * d;
-
-  auto block_max = max_data[blockIdx.x];
-  int64_t step = BlockDim * remain;
-
-  // In numeric stable mode softmax_with_loss, we calc loss with
-  // tmp_i_j = x_i_j - max_i - logDiffMaxSum_i, instead of
-  // log(exp(x_i_j - max_i)/DiffMaxSum_i). Therefore, log(0) will not occur.
-  // Also we calc softmax_i_j = e^{tmp_i_j}, the maximum and minimum value will
-  // be 1.0 and 0.0, represent prob is 1.0 and 0.0.
-  // So there is no need to clip on shift_softmax.
-  softmax[beg_idx] = logits_data[beg_idx] - block_max;
-  T diff_max_sum = exp_on_device(softmax[beg_idx]);
-  auto idx = beg_idx + step;
-  while (idx < end_idx) {
-    softmax[idx] = logits_data[idx] - block_max;
-    diff_max_sum += exp_on_device(softmax[idx]);
-    idx += step;
-  }
-
-  diff_max_sum =
-      BlockReduce<T, BlockDim>(temp_storage).Reduce(diff_max_sum, cub::Sum());
-  if (threadIdx.x == 0) max_data[blockIdx.x] = log_on_device(diff_max_sum);
-
-  if (!CalculateLogSoftmax) return;
-  __syncthreads();
-  diff_max_sum = max_data[blockIdx.x];
-  softmax[beg_idx] -= diff_max_sum;
-  beg_idx += step;
-  while (beg_idx < end_idx) {
-    softmax[beg_idx] -= diff_max_sum;
-    beg_idx += step;
-  }
-
-  // Note(zhiqiu): since different threads may use max_data[blockIdx.x] to
-  // calculate diff_max_sum, __syncthreads() is needed here.
-  __syncthreads();
-  if (threadIdx.x == 0) max_data[blockIdx.x] = 0;
-}
-
-#ifdef __HIPCC__  // @{ HIP Seperate Kernel for RowReductionForDiffMaxSum
-// Note(qili93): HIP do not support return in kernel, need to seperate
-// RowReductionForDiffMaxSum into two kernels below
-template <typename T, int BlockDim>
-static __global__ void RowReductionForSum(const T* logits_data, T* max_data,
-                                          T* softmax, int64_t d, int axis_dim) {
-  __shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
-
-  int64_t remain = d / axis_dim;
-  int64_t idx_n = blockIdx.x / remain;
-  int64_t idx_remain = blockIdx.x % remain;
-  int64_t beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int64_t end_idx = (idx_n + 1) * d;
-
-  auto block_max = max_data[blockIdx.x];
-  int64_t step = BlockDim * remain;
-
-  softmax[beg_idx] = logits_data[beg_idx] - block_max;
-  T diff_max_sum = exp_on_device(softmax[beg_idx]);
-  auto idx = beg_idx + step;
-  while (idx < end_idx) {
-    softmax[idx] = logits_data[idx] - block_max;
-    diff_max_sum += exp_on_device(softmax[idx]);
-    idx += step;
-  }
-
-  diff_max_sum =
-      BlockReduce<T, BlockDim>(temp_storage).Reduce(diff_max_sum, cub::Sum());
-  if (threadIdx.x == 0) max_data[blockIdx.x] = log_on_device(diff_max_sum);
-}
-
-template <typename T, int BlockDim, bool CalculateLogSoftmax = false>
-static __global__ void RowReductionForDiff(const T* logits_data, T* max_data,
-                                           T* softmax, int d, int axis_dim) {
-  int remain = d / axis_dim;
-  int idx_n = blockIdx.x / remain;
-  int idx_remain = blockIdx.x % remain;
-  int beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int end_idx = (idx_n + 1) * d;
-  int step = BlockDim * remain;
-
-  T diff_max_sum = max_data[blockIdx.x];
-  softmax[beg_idx] -= diff_max_sum;
-  beg_idx += step;
-  while (beg_idx < end_idx) {
-    softmax[beg_idx] -= diff_max_sum;
-    beg_idx += step;
-  }
-
-  __syncthreads();
-  if (threadIdx.x == 0) max_data[blockIdx.x] = 0;
-}
-#endif  // @} End HIP Seperate Kernel for RowReductionForDiffMaxSum
-
-// Make sure that BlockDim <= axis_dim
-template <typename T, int BlockDim>
-static __global__ void RowReductionForSoftmaxAndCrossEntropy(
-    const T* logits_data, const T* labels_data, T* loss_data, T* softmax,
-    int64_t d, int axis_dim) {
-  __shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
-
-  // logits, softmax, labels data view as [n, axis_dim, remain]
-  // loss_data view as [n, 1, remain]
-  // blockDim = n * remain, split blockIdx to idx_n and idx_remain
-  int64_t remain = d / axis_dim;
-  int64_t idx_n = blockIdx.x / remain;
-  int64_t idx_remain = blockIdx.x % remain;
-  int64_t beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int64_t end_idx = (idx_n + 1) * d;
-
-  // log_diff_max_sum shares memory with loss
-  auto block_log_diff_max_sum = loss_data[blockIdx.x];
-  auto tmp = softmax[beg_idx] - block_log_diff_max_sum;
-  softmax[beg_idx] = exp_on_device(tmp);
-  auto loss = -labels_data[beg_idx] * tmp;
-  int64_t step = BlockDim * remain;
-  beg_idx += step;
-  while (beg_idx < end_idx) {
-    tmp = softmax[beg_idx] - block_log_diff_max_sum;
-    softmax[beg_idx] = exp_on_device(tmp);
-    loss -= (labels_data[beg_idx] * tmp);
-    beg_idx += step;
-  }
-
-  loss = BlockReduce<T, BlockDim>(temp_storage).Reduce(loss, cub::Sum());
-  if (threadIdx.x == 0) loss_data[blockIdx.x] = loss;
-}
-
-// Make sure that BlockDim <= axis_dim
-template <typename T, int BlockDim>
-static __global__ void RowReductionForCrossEntropy(const T* logits_data,
-                                                   const T* labels_data,
-                                                   T* loss_data, int d,
-                                                   int axis_dim) {
-  __shared__ BlockReduceTempStorage<T, BlockDim> temp_storage;
-
-  // logits, softmax, labels data view as [n, axis_dim, remain]
-  // loss_data view as [n, 1, remain]
-  // blockDim = n * remain, split blockIdx to idx_n and idx_remain
-  int remain = d / axis_dim;
-  int idx_n = blockIdx.x / remain;
-  int idx_remain = blockIdx.x % remain;
-  int beg_idx = idx_n * d + threadIdx.x * remain + idx_remain;
-  int end_idx = (idx_n + 1) * d;
-
-  // log_diff_max_sum shares memory with loss
-  auto block_log_diff_max_sum = loss_data[blockIdx.x];
-  auto tmp = log_on_device(logits_data[beg_idx]);  // when not with softmax,
-                                                   // softmax is stored in
-                                                   // logits_data
-  auto loss = -labels_data[beg_idx] * tmp;
-  int step = BlockDim * remain;
-  beg_idx += step;
-  while (beg_idx < end_idx) {
-    tmp = log_on_device(logits_data[beg_idx]);  // when not with softmax,
-                                                // softmax is stored in
-                                                // logits_data
-    loss -= (labels_data[beg_idx] * tmp);
-    beg_idx += step;
-  }
-
-  loss = BlockReduce<T, BlockDim>(temp_storage).Reduce(loss, cub::Sum());
-  if (threadIdx.x == 0) loss_data[blockIdx.x] = loss;
-}
-
-template <typename T>
-static void SoftmaxWithCrossEntropyFusedKernel(
-    const T* logits_data, const T* labels_data, T* softmax_data, T* loss_data,
-    int64_t n, int64_t d, int axis_dim, gpuStream_t stream) {
-#ifdef __HIPCC__
-  constexpr int kMaxBlockDim = 256;
-#else
-  constexpr int kMaxBlockDim = 512;
-#endif
-  int64_t block_dim = axis_dim >= kMaxBlockDim
-                          ? kMaxBlockDim
-                          : (1 << static_cast<int>(std::log2(axis_dim)));
-  int64_t grid_dim = n * d / axis_dim;
-#ifdef __HIPCC__
-#define CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(BlockDim)                 \
-  case BlockDim:                                                               \
-    hipLaunchKernelGGL(HIP_KERNEL_NAME(RowReductionForMax<T, BlockDim>),       \
-                       dim3(grid_dim), dim3(BlockDim), 0, stream, logits_data, \
-                       loss_data, d, axis_dim);                                \
-    hipLaunchKernelGGL(HIP_KERNEL_NAME(RowReductionForSum<T, BlockDim>),       \
-                       dim3(grid_dim), dim3(BlockDim), 0, stream, logits_data, \
-                       loss_data, softmax_data, d, axis_dim);                  \
-    hipLaunchKernelGGL(                                                        \
-        HIP_KERNEL_NAME(RowReductionForSoftmaxAndCrossEntropy<T, BlockDim>),   \
-        dim3(grid_dim), dim3(BlockDim), 0, stream, logits_data, labels_data,   \
-        loss_data, softmax_data, d, axis_dim);                                 \
-    break
-#else
-#define CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(BlockDim)                 \
-  case BlockDim:                                                               \
-    RowReductionForMax<T, BlockDim><<<grid_dim, BlockDim, 0, stream>>>(        \
-        logits_data, loss_data, d, axis_dim);                                  \
-    RowReductionForDiffMaxSum<T, BlockDim><<<grid_dim, BlockDim, 0, stream>>>( \
-        logits_data, loss_data, softmax_data, d, axis_dim);                    \
-    RowReductionForSoftmaxAndCrossEntropy<                                     \
-        T, BlockDim><<<grid_dim, BlockDim, 0, stream>>>(                       \
-        logits_data, labels_data, loss_data, softmax_data, d, axis_dim);       \
-    break
-#endif
-
-  switch (block_dim) {
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(512);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(256);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(128);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(64);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(32);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(16);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(8);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(4);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(2);
-    default:
-      PADDLE_THROW(platform::errors::Unavailable(
-          "Block Dimension must be 2^n in softmax_with_cross_entropy_op."));
-      break;
-  }
-
-#undef CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL
-}
-
-// not with softmax
-template <typename T>
-static void CrossEntropyFusedKernel(const T* logits_data, const T* labels_data,
-                                    T* loss_data, int n, int d, int axis_dim,
-                                    gpuStream_t stream) {
-  constexpr int kMaxBlockDim = 512;
-  int block_dim = axis_dim >= kMaxBlockDim
-                      ? kMaxBlockDim
-                      : (1 << static_cast<int>(std::log2(axis_dim)));
-  int grid_dim = n * d / axis_dim;
-
-#define CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(BlockDim)                \
-  case BlockDim:                                                              \
-    RowReductionForCrossEntropy<T,                                            \
-                                BlockDim><<<grid_dim, BlockDim, 0, stream>>>( \
-        logits_data, labels_data, loss_data, d, axis_dim);                    \
-    break
-
-  switch (block_dim) {
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(512);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(256);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(128);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(64);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(32);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(16);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(8);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(4);
-    CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL(2);
-    default:
-      PADDLE_THROW(platform::errors::Unavailable(
-          "Block Dimension must be 2^n in softmax_with_cross_entropy_op."));
-      break;
-  }
-
-#undef CALL_SOFTMAX_WITH_CROSS_ENTROPY_FUSED_KERNEL
-}
-
 template <typename T>
 class SoftmaxWithCrossEntropyCUDAKernel : public framework::OpKernel<T> {
  public:

  using AccT = typename details::MPTypeTrait<T>::Type;

  // use 128 threads per block to maximimize gpu utilization
  const int Log2Elements = static_cast<int>(Log2Ceil(element_count));
  const int kDimCeil = 1 << Log2Elements;
  int kWarpSize = (kDimCeil < 32) ? kDimCeil : 32;
  int batches_per_warp = (kDimCeil <= 128) ? 2 : 1;
  constexpr int threads_per_block = 128;

    if (soft_label) {
      auto* logits_data = logits->data<T>();
      auto* labels_data = labels->data<T>();
      SoftmaxWithCrossEntropySoftLabel<T>(
          context.cuda_device_context(), rank, axis, logits_data, labels_data,
          softmax_data, loss_data, n, axis_dim, d / axis_dim);
    } else {
      if (!context.Attr<bool>("numeric_stable_mode")) {
        // CUDNN kernel only suppoer 2-D tensor and perfome softmax on last dim
