diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp
index fb340cb43..396f1f915 100644
--- a/src/infinicore/pybind11/ops.hpp
+++ b/src/infinicore/pybind11/ops.hpp
@@ -113,8 +113,8 @@
 #include "ops/rwkv5_wkv.hpp"
 #include "ops/scal.hpp"
 #include "ops/scatter.hpp"
-#include "ops/selu.hpp"
 #include "ops/select_last_token_hidden.hpp"
+#include "ops/selu.hpp"
 #include "ops/sigmoid.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"
diff --git a/src/infiniop/ops/softmax/metax/softmax_metax.maca b/src/infiniop/ops/softmax/metax/softmax_metax.maca
index f778440b6..6aef196c0 100644
--- a/src/infiniop/ops/softmax/metax/softmax_metax.maca
+++ b/src/infiniop/ops/softmax/metax/softmax_metax.maca
@@ -1,187 +1,164 @@
 #include "../../../devices/metax/metax_common.h"
-#include "../../../devices/metax/metax_handle.h"
-#include "../../../devices/metax/metax_kernel_common.h"
 #include "softmax_metax.h"
-#include <cmath>
-#include <limits>
-
-namespace op::softmax::metax {
-namespace {
-
-__device__ __forceinline__ float to_float(float val) { return val; }
-__device__ __forceinline__ float to_float(double val) { return static_cast<float>(val); }
-__device__ __forceinline__ float to_float(__half val) { return __half2float(val); }
-__device__ __forceinline__ float to_float(cuda_bfloat16 val) { return __bfloat162float(val); }
-
-template <typename T>
-__device__ __forceinline__ T from_float(float val) { return static_cast<T>(val); }
-
-template <>
-__device__ __forceinline__ __half from_float<__half>(float val) { return __float2half(val); }
 
-template <>
-__device__ __forceinline__ cuda_bfloat16 from_float<cuda_bfloat16>(float val) { return __float2bfloat16(val); }
-
-template <typename T>
-__device__ __forceinline__ T warp_reduce_max(T val) {
-    for (int offset = 16; offset > 0; offset /= 2) {
-        T shuffled = __shfl_down_sync(0xffffffff, val, offset);
-        val = val > shuffled ? val : shuffled;
-    }
-    return val;
-}
-
-template <typename T>
-__device__ __forceinline__ T warp_reduce_sum(T val) {
-    for (int offset = 16; offset > 0; offset /= 2) {
-        val += __shfl_down_sync(0xffffffff, val, offset);
-    }
-    return val;
-}
+#ifdef ENABLE_METAX_MC_API
+#include <cub/block/block_reduce.cuh>
+#else
+#include <hccub/block/block_reduce.cuh>
+#endif
+#include "../../../devices/metax/metax_kernel_common.h"
 
-__device__ __forceinline__ float block_reduce_max(float val) {
-    __shared__ float shared[32];
-    int lane = threadIdx.x & 31;
-    int wid = threadIdx.x >> 5;
-    val = warp_reduce_max(val);
-    if (lane == 0) {
-        shared[wid] = val;
-    }
-    __syncthreads();
-    val = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? shared[lane] : -INFINITY;
-    if (wid == 0) {
-        val = warp_reduce_max(val);
-    }
-    return val;
-}
+#include "../cuda/kernel.cuh"
 
-__device__ __forceinline__ float block_reduce_sum(float val) {
-    __shared__ float shared[32];
-    int lane = threadIdx.x & 31;
-    int wid = threadIdx.x >> 5;
-    val = warp_reduce_sum(val);
-    if (lane == 0) {
-        shared[wid] = val;
-    }
-    __syncthreads();
-    val = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? shared[lane] : 0.0f;
-    if (wid == 0) {
-        val = warp_reduce_sum(val);
-    }
-    return val;
+template <typename Tdata, unsigned int BLOCK_SIZE>
+INFINIOP_METAX_KERNEL blockSoftmax(
+    Tdata *y, const Tdata *x,
+    size_t dimsize,
+    ptrdiff_t stride) {
+    blockSoftmaxKernel<Tdata, BLOCK_SIZE>(x, y, dimsize, stride);
 }
 
-template <typename T>
-__global__ void softmax_kernel(
-    T *__restrict__ output,
-    const T *__restrict__ input,
+template <typename Tdata, unsigned int BLOCK_SIZE_x, unsigned int BLOCK_SIZE_y, int numPerThreadx>
+INFINIOP_METAX_KERNEL warpSoftmax(
+    Tdata *y, const Tdata *x,
     size_t othersize,
     size_t dimsize,
     ptrdiff_t stride) {
-    const size_t other_idx = static_cast<size_t>(blockIdx.x);
-    if (other_idx >= othersize) {
-        return;
-    }
-
-    const size_t inner_idx = other_idx % static_cast<size_t>(stride);
-    const size_t outer_idx = other_idx / static_cast<size_t>(stride);
-    const size_t base_offset = outer_idx * dimsize * static_cast<size_t>(stride) + inner_idx;
-
-    float local_max = -INFINITY;
-    for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) {
-        const float val = to_float(input[base_offset + i * stride]);
-        local_max = local_max > val ? local_max : val;
-    }
-
-    __shared__ float s_max;
-    __shared__ float s_sum;
-    const float max_val = block_reduce_max(local_max);
-    if (threadIdx.x == 0) {
-        s_max = max_val;
-    }
-    __syncthreads();
-
-    float local_sum = 0.0f;
-    for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) {
-        local_sum += expf(to_float(input[base_offset + i * stride]) - s_max);
-    }
-
-    const float sum_val = block_reduce_sum(local_sum);
-    if (threadIdx.x == 0) {
-        s_sum = sum_val;
-    }
-    __syncthreads();
-
-    const float inv_sum = 1.0f / s_sum;
-    for (size_t i = threadIdx.x; i < dimsize; i += blockDim.x) {
-        const size_t idx = base_offset + i * stride;
-        output[idx] = from_float<T>(expf(to_float(input[idx]) - s_max) * inv_sum);
-    }
+    warpSoftmaxKernel<Tdata, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>(x, y, othersize, dimsize, stride);
 }
 
-template <typename T>
-void launch_kernel(void *output, const void *input, const SoftmaxInfo &info, void *stream) {
-    auto hc_stream = reinterpret_cast<hcStream_t>(stream);
-    unsigned int threads_per_block = 256;
-    if (info.dimsize < 256) {
-        threads_per_block = 128;
-    }
-    if (info.dimsize < 128) {
-        threads_per_block = 64;
-    }
-    if (info.dimsize < 64) {
-        threads_per_block = 32;
-    }
-    softmax_kernel<T><<<info.othersize, threads_per_block, 0, hc_stream>>>(
-        reinterpret_cast<T *>(output),
-        reinterpret_cast<const T *>(input),
-        info.othersize,
-        info.dimsize,
-        info.stride);
-}
-
-} // namespace
+namespace op::softmax::metax {
 
-struct Descriptor::Opaque {};
+struct Descriptor::Opaque {
+    std::shared_ptr<device::metax::Handle::Internal> internal;
+};
 
 Descriptor::~Descriptor() {
     delete _opaque;
 }
 
 infiniStatus_t Descriptor::create(
-    infiniopHandle_t handle_,
+    infiniopHandle_t handle,
     Descriptor **desc_ptr,
     infiniopTensorDescriptor_t y_desc,
     infiniopTensorDescriptor_t x_desc,
     int axis) {
-    auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
     auto info = SoftmaxInfo::create(y_desc, x_desc, axis);
     CHECK_RESULT(info);
-    *desc_ptr = new Descriptor(new Opaque(), info.take(), 0, handle->device, handle->device_id);
+    *desc_ptr = new Descriptor(
+        new Opaque{reinterpret_cast<device::metax::Handle *>(handle)->internal()},
+        info.take(), 0, handle->device, handle->device_id);
     return INFINI_STATUS_SUCCESS;
 }
 
-infiniStatus_t Descriptor::calculate(
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    void *stream) const {
-    (void)workspace;
-    (void)workspace_size;
-    switch (_info.dtype) {
-    case INFINI_DTYPE_F16:
-        launch_kernel<__half>(y, x, _info, stream);
-        break;
-    case INFINI_DTYPE_BF16:
-        launch_kernel<cuda_bfloat16>(y, x, _info, stream);
-        break;
-    case INFINI_DTYPE_F32:
-        launch_kernel<float>(y, x, _info, stream);
-        break;
-    default:
+template <unsigned int BLOCK_SIZE>
+infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype,
+                            size_t othersize, size_t dimsize, ptrdiff_t stride,
+                            hcStream_t stream) {
+    int num_blocks = (int)othersize;
+    // Avoid the warp kernel register pressure on large, non-contiguous rows.
+    const bool use_block_kernel = dimsize > 1024 || (dimsize >= 512 && stride != 1);
+    if (dtype == INFINI_DTYPE_F16) {
+        if (use_block_kernel) {
+            blockSoftmax<half, BLOCK_SIZE>
+                <<<num_blocks, BLOCK_SIZE, 0, stream>>>((half *)y, (const half *)x,
+                                                        dimsize, stride);
+        } else if (dimsize > 31) {
+            constexpr unsigned int BLOCK_SIZE_x = 32;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 32;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<half, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((half *)y, (const half *)x,
+                                                     othersize, dimsize, stride);
+        } else {
+            constexpr unsigned int BLOCK_SIZE_x = 16;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 2;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<half, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((half *)y, (const half *)x,
+                                                     othersize, dimsize, stride);
+        }
+
+    } else if (dtype == INFINI_DTYPE_BF16) {
+        if (use_block_kernel) {
+            blockSoftmax<cuda_bfloat16, BLOCK_SIZE>
+                <<<num_blocks, BLOCK_SIZE, 0, stream>>>((cuda_bfloat16 *)y, (const cuda_bfloat16 *)x,
+                                                        dimsize, stride);
+        } else if (dimsize > 31) {
+            constexpr unsigned int BLOCK_SIZE_x = 32;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 32;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<cuda_bfloat16, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((cuda_bfloat16 *)y, (const cuda_bfloat16 *)x,
+                                                     othersize, dimsize, stride);
+        } else {
+            constexpr unsigned int BLOCK_SIZE_x = 16;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 2;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<cuda_bfloat16, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((cuda_bfloat16 *)y, (const cuda_bfloat16 *)x,
+                                                     othersize, dimsize, stride);
+        }
+
+    } else if (dtype == INFINI_DTYPE_F32) {
+        if (use_block_kernel) {
+            blockSoftmax<float, BLOCK_SIZE>
+                <<<num_blocks, BLOCK_SIZE, 0, stream>>>((float *)y, (const float *)x,
+                                                        dimsize, stride);
+        } else if (dimsize > 31) {
+            constexpr unsigned int BLOCK_SIZE_x = 32;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 32;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<float, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((float *)y, (const float *)x,
+                                                     othersize, dimsize, stride);
+        } else {
+            constexpr unsigned int BLOCK_SIZE_x = 16;
+            constexpr unsigned int BLOCK_SIZE_y = BLOCK_SIZE / 32;
+            constexpr int numPerThreadx = 2;
+            int num_block_x = (num_blocks + BLOCK_SIZE_y - 1) / BLOCK_SIZE_y;
+            dim3 block_dim(BLOCK_SIZE_x, BLOCK_SIZE_y, 1);
+            dim3 grid_dim(num_block_x, 1, 1);
+            warpSoftmax<float, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+                <<<grid_dim, block_dim, 0, stream>>>((float *)y, (const float *)x,
+                                                     othersize, dimsize, stride);
+        }
+    } else {
         return INFINI_STATUS_BAD_TENSOR_DTYPE;
     }
+    CHECK_METAX(hcGetLastError());
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size,
+                                     void *y,
+                                     const void *x,
+                                     void *stream_) const {
+    hcStream_t stream = (hcStream_t)stream_;
+    if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) {
+        CHECK_STATUS(launchKernel<METAX_BLOCK_SIZE_1024>(
+            y, x, _info.dtype, _info.othersize, _info.dimsize, _info.stride, stream));
+    } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) {
+        CHECK_STATUS(launchKernel<METAX_BLOCK_SIZE_512>(
+            y, x, _info.dtype, _info.othersize, _info.dimsize, _info.stride, stream));
+    } else {
+        return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
+    }
     return INFINI_STATUS_SUCCESS;
 }