     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 == 128 && N % warps_per_block == 0 && sizeof(T) <= 4) {
-      // 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");
+    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;
+        }
       }
-    } else {
+    }
+    if (!optimize) {
       ScopedTensorDescriptor desc;
       std::vector<int> tensor_dims = {N, dim, D, 1};
       DataLayout layout = DataLayout::kNCHW;

#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 {

using DataLayout = platform::DataLayout;
using Tensor = framework::Tensor;

#define LAUNCH_SOFTMAX_WARP_FORWARD(Log2Elements)                  \
  case Log2Elements:                                               \
    WarpSoftmaxForward<T, float, Log2Elements><<<                  \
        blocks, threads, 0, ctx.cuda_device_context().stream()>>>( \
        out_data, x->data<T>(), N, dim, dim);                      \
    break;

static inline int SizeOutAxis(const int axis, DDim dims) {
  int size = 1;
  for (int i = axis + 1; i < dims.size(); i++) {
    size *= dims[i];
  }
  return size;
}

int log2_ceil(int value) {
  int log2_value = 0;
  while ((1 << log2_value) < value) ++log2_value;
  return log2_value;
}

template <typename T, int VLEN>
