   reinterpret_cast<VECT*>(&dst[offset + idx])[0] = buf;
 }
 
+template <typename T, int WARP_BATCH, int WARP_SIZE_SOFTMAX>
+__device__ __forceinline__ void warp_reduce_sum(T* sum) {
+#pragma unroll
+  for (int offset = WARP_SIZE_SOFTMAX / 2; offset > 0; offset /= 2) {
+#pragma unroll
+    for (int i = 0; i < WARP_BATCH; ++i) {
+      T sum_val = platform::CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
+      sum[i] = sum[i] + sum_val;
+    }
+  }
+}
+
+template <typename T, int WARP_BATCH, int WARP_SIZE_SOFTMAX>
+__device__ __forceinline__ void warp_reduce_max(T* sum) {
+#pragma unroll
+  for (int offset = WARP_SIZE_SOFTMAX / 2; offset > 0; offset /= 2) {
+#pragma unroll
+    for (int i = 0; i < WARP_BATCH; ++i) {
+      T max_val = platform::CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
+      sum[i] = max(sum[i], max_val);
+    }
+  }
+}
+
+template <typename T, typename AccT, int Log2Elements>
+__global__ void WarpSoftmaxForward(T* dst, const T* src, const int batch_size,
+                                   const int stride, const int element_count) {
+  constexpr int next_power_of_two = 1 << Log2Elements;
+  constexpr int warp_size_softmax =
+      (next_power_of_two < 32) ? next_power_of_two : 32;
+  constexpr int WARP_ITERATIONS = next_power_of_two / warp_size_softmax;
+  constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+
+  int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+
+  int local_batches = batch_size - first_batch;
+  if (local_batches > WARP_BATCH) {
+    local_batches = WARP_BATCH;
+  }
+
+  int local_idx = threadIdx.x;
+
+  src += first_batch * stride + local_idx;
+  dst += first_batch * stride + local_idx;
+
+  // load data from global memory
+  AccT elements[WARP_BATCH][WARP_ITERATIONS];
+  for (int i = 0; i < WARP_BATCH; ++i) {
+    int batch_element_count = (i >= local_batches) ? 0 : element_count;
+    for (int it = 0; it < WARP_ITERATIONS; ++it) {
+      int element_index = local_idx + it * warp_size_softmax;
+      if (element_index < batch_element_count) {
+        elements[i][it] =
+            static_cast<float>(src[i * element_count + it * warp_size_softmax]);
+      } else {
+        elements[i][it] = -std::numeric_limits<AccT>::infinity();
+      }
+    }
+  }
+
+  // compute max_value
+  AccT max_value[WARP_BATCH];
+#pragma unroll
+  for (int i = 0; i < WARP_BATCH; ++i) {
+    max_value[i] = elements[i][0];
+#pragma unroll
+    for (int it = 1; it < WARP_ITERATIONS; ++it) {
+      max_value[i] =
+          (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+    }
+  }
+  warp_reduce_max<AccT, WARP_BATCH, warp_size_softmax>(max_value);
+
+  AccT sum[WARP_BATCH]{0.0f};
+#pragma unroll
+  for (int i = 0; i < WARP_BATCH; ++i) {
+#pragma unroll
+    for (int it = 0; it < WARP_ITERATIONS; ++it) {
+      elements[i][it] = (std::exp((elements[i][it] - max_value[i])));
+      sum[i] += elements[i][it];
+    }
+  }
+  warp_reduce_sum<AccT, WARP_BATCH, warp_size_softmax>(sum);
+
+// store result
+#pragma unroll
+  for (int i = 0; i < WARP_BATCH; ++i) {
+    if (i >= local_batches) break;
+#pragma unroll
+    for (int it = 0; it < WARP_ITERATIONS; ++it) {
+      int element_index = local_idx + it * warp_size_softmax;
+      if (element_index < element_count) {
+        dst[i * element_count + it * warp_size_softmax] =
+            elements[i][it] / sum[i];
+      } else {
+        break;
+      }
+    }
+  }
+}
+
 template <typename T, int VPT, int WARP_PER_BLOCK>
 __global__ void VecSoftmaxBackward(T* dst, const T* grad, const T* src,
                                    const int batch_size,

#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/math/math_cuda_utils.h"
#include "paddle/fluid/operators/softmax_op.h"
#include "paddle/fluid/platform/cudnn_helper.h"

namespace paddle {

    const int N = SizeToAxis(axis, dims);
    const int D = SizeOutAxis(axis, dims);

    constexpr int max_dim = 320;
    bool optimize = false;
    constexpr int warps_per_block = 4;
    if (D == 1 && dim <= max_dim && sizeof(T) <= 4) {
      if (dim == 128 && N % warps_per_block == 0) {
        optimize = true;
        // a warp for a batch, 4 elements for a thread, only support the softmax
        // dim size = 128 currently
        if (sizeof(T) == 2) {
          VecSoftmaxForward<T, int2, 4, warps_per_block><<<
              N / warps_per_block, warps_per_block * WARP_SIZE, 0,
              ctx.cuda_device_context().stream()>>>(out_data, x->data<T>(), N,
                                                    dim);
        } else if (sizeof(T) == 4) {
          VecSoftmaxForward<T, int4, 4, warps_per_block><<<
              N / warps_per_block, warps_per_block * WARP_SIZE, 0,
              ctx.cuda_device_context().stream()>>>(out_data, x->data<T>(), N,
                                                    dim);
        } else {
          assert(false && "not support");
        }
      } else if (dim < max_dim) {
        optimize = true;
        int log2_elements = static_cast<int>(log2_ceil(dim));
        const int next_power_of_two = 1 << log2_elements;

        int warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;

        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;

        // use 128 threads per block to maximimize gpu utilization
        constexpr int threads_per_block = 128;

        int warps_per_block = (threads_per_block / warp_size);
        int batches_per_block = warps_per_block * batches_per_warp;
        int blocks = (N + batches_per_block - 1) / batches_per_block;
        dim3 threads(warp_size, warps_per_block, 1);

        switch (log2_elements) {
          LAUNCH_SOFTMAX_WARP_FORWARD(0);  // 1
          LAUNCH_SOFTMAX_WARP_FORWARD(1);  // 2
          LAUNCH_SOFTMAX_WARP_FORWARD(2);  // 4
          LAUNCH_SOFTMAX_WARP_FORWARD(3);  // 8
          LAUNCH_SOFTMAX_WARP_FORWARD(4);  // 16
          LAUNCH_SOFTMAX_WARP_FORWARD(5);  // 32
          LAUNCH_SOFTMAX_WARP_FORWARD(6);  // 64
          LAUNCH_SOFTMAX_WARP_FORWARD(7);  // 128
          LAUNCH_SOFTMAX_WARP_FORWARD(8);  // 256
          LAUNCH_SOFTMAX_WARP_FORWARD(9);  // 512
          default:
            break;
        }
      }
    }
    if (!optimize) {
      ScopedTensorDescriptor desc;
      std::vector<int> tensor_dims = {N, dim, D, 1};
      DataLayout layout = DataLayout::kNCHW;
