Skip to content

feat: support softmax operator on metax #1347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wooway777 merged 2 commits into from
Jul 29, 2026
Merged

feat: support softmax operator on metax #1347

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dd1824a
feat: support softmax operator on metax
LindseyMei Jun 29, 2026
f0e64c4
Merge main and optimize MetaX softmax dispatch
wooway777 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/infinicore/pybind11/ops.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
285 changes: 131 additions & 154 deletions src/infiniop/ops/softmax/metax/softmax_metax.maca
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
}

Expand Down
Loading