diff --git a/src/infiniop/ops/quickgelu/cuda/kernel.cuh b/src/infiniop/ops/quickgelu/cuda/kernel.cuh index 2c13c4b9d..4042fe3bb 100644 --- a/src/infiniop/ops/quickgelu/cuda/kernel.cuh +++ b/src/infiniop/ops/quickgelu/cuda/kernel.cuh @@ -1,58 +1,31 @@ #ifndef __QUICKGELU_CUDA_H__ #define __QUICKGELU_CUDA_H__ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" -#include -#include +#include namespace op::quickgelu::cuda { typedef struct QuickGeluOp { public: static constexpr size_t num_inputs = 1; - template __device__ __forceinline__ T operator()(const T &x) const { - // quickgelu(x) = x * sigmoid(1.702 * x) - + // quickgelu(x) = x * sigmoid(1.702 * x) = x / (1 + exp(-1.702 x)) constexpr float alpha = 1.702f; - - if constexpr (std::is_same_v) { - half2 ax = __hmul2(make_half2(alpha, alpha), x); - half2 denominator = __hadd2(make_half2(1, 1), h2exp(__hneg2(ax))); - half2 sigmoid = h2rcp(denominator); - return __hmul2(x, sigmoid); - + if constexpr (std::is_same_v) { + float x_f = __bfloat162float(x); + float result = x_f / (1.0f + expf(-alpha * x_f)); + return __float2bfloat16(result); } else if constexpr (std::is_same_v) { - half ax = __hmul(__float2half(alpha), x); - half denominator = __hadd(__float2half(1.0f), hexp(__hneg(ax))); - half sigmoid = hrcp(denominator); - return __hmul(x, sigmoid); - - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(x); - float ax = alpha * xf; - float s = 1.0f / (1.0f + __expf(-ax)); - return __float2bfloat16(xf * s); - + float x_f = __half2float(x); + float result = x_f / (1.0f + expf(-alpha * x_f)); + return __float2half(result); } else if constexpr (std::is_same_v) { - float ax = alpha * x; - float s; - if (ax >= 0.0f) { - float z = expf(-ax); - s = 1.0f / (1.0f + z); - } else { - float z = expf(ax); - s = z / (1.0f + z); - } - return x * s; - - } else { // double - double ax = static_cast(alpha) * x; - return x / (1.0 + exp(-ax)); + return x / (1.0f + expf(-alpha * x)); + } else { + return x / (1.0 + exp(-static_cast(alpha) * x)); } } - } QuickGeluOp; } // namespace op::quickgelu::cuda diff --git a/src/infiniop/ops/quickgelu/metax/quickgelu_metax.h b/src/infiniop/ops/quickgelu/metax/quickgelu_metax.h new file mode 100644 index 000000000..9596ddf0e --- /dev/null +++ b/src/infiniop/ops/quickgelu/metax/quickgelu_metax.h @@ -0,0 +1,8 @@ +#ifndef __QUICKGELU_METAX_API_H__ +#define __QUICKGELU_METAX_API_H__ + +#include "../../../elementwise/metax/elementwise_metax_api.h" + +ELEMENTWISE_DESCRIPTOR(quickgelu, metax) + +#endif // __QUICKGELU_METAX_API_H__ diff --git a/src/infiniop/ops/quickgelu/metax/quickgelu_metax.maca b/src/infiniop/ops/quickgelu/metax/quickgelu_metax.maca new file mode 100644 index 000000000..b85472505 --- /dev/null +++ b/src/infiniop/ops/quickgelu/metax/quickgelu_metax.maca @@ -0,0 +1,65 @@ +#include "quickgelu_metax.h" + +#include "../../../elementwise/metax/elementwise_metax.h" + +#include "../cuda/kernel.cuh" + +namespace op::quickgelu::metax { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + + const auto &input_desc = input_desc_vec.at(0); + const auto &output_shape = out_desc->shape(); + const auto &input_shape = input_desc->shape(); + + CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); + + CHECK_SAME_SHAPE(output_shape, input_shape); + + // create METAX elementwise descriptor + CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + // NOTE(ernie-vl / metax): QuickGeluOp (shared cuda/kernel.cuh) selects its bf16 + // path via std::is_same_v. On metax both cuda_bfloat16 and + // __nv_bfloat16 alias the same maca bf16 type, so instantiating with + // cuda_bfloat16 (as gelu does) hits the bf16 branch. VERIFY(C500): if the bf16 + // branch is missed at compile time, change QuickGeluOp's check to cuda_bfloat16. + switch (_dtype) { + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, cuda::QuickGeluOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F16: + return _device_info->calculate<256, cuda::QuickGeluOp, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, cuda::QuickGeluOp, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, cuda::QuickGeluOp, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::quickgelu::metax diff --git a/src/infiniop/ops/quickgelu/nvidia/quickgelu_nvidia.cu b/src/infiniop/ops/quickgelu/nvidia/quickgelu_nvidia.cu index 387e08ecb..525ff9f67 100644 --- a/src/infiniop/ops/quickgelu/nvidia/quickgelu_nvidia.cu +++ b/src/infiniop/ops/quickgelu/nvidia/quickgelu_nvidia.cu @@ -1,5 +1,6 @@ -#include "../cuda/kernel.cuh" #include "quickgelu_nvidia.cuh" +#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" +#include "../cuda/kernel.cuh" namespace op::quickgelu::nvidia { diff --git a/src/infiniop/ops/quickgelu/operator.cc b/src/infiniop/ops/quickgelu/operator.cc index f85a3e49a..851c96530 100644 --- a/src/infiniop/ops/quickgelu/operator.cc +++ b/src/infiniop/ops/quickgelu/operator.cc @@ -8,6 +8,9 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) #include "nvidia/quickgelu_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/quickgelu_metax.h" +#endif __INFINI_C infiniStatus_t infiniopCreateQuickGeluDescriptor( infiniopHandle_t handle, @@ -38,6 +41,9 @@ __INFINI_C infiniStatus_t infiniopCreateQuickGeluDescriptor( #endif #ifdef ENABLE_HYGON_API CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -68,6 +74,9 @@ __INFINI_C infiniStatus_t infiniopGetQuickGeluWorkspaceSize(infiniopQuickGeluDes #endif #ifdef ENABLE_HYGON_API GET(INFINI_DEVICE_HYGON, nvidia) +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -104,6 +113,9 @@ __INFINI_C infiniStatus_t infiniopQuickGelu( #endif #ifdef ENABLE_HYGON_API CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -134,6 +146,9 @@ __INFINI_C infiniStatus_t infiniopDestroyQuickGeluDescriptor(infiniopQuickGeluDe #endif #ifdef ENABLE_HYGON_API DELETE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/softmax/metax/softmax_metax.h b/src/infiniop/ops/softmax/metax/softmax_metax.h new file mode 100644 index 000000000..d05575647 --- /dev/null +++ b/src/infiniop/ops/softmax/metax/softmax_metax.h @@ -0,0 +1,8 @@ +#ifndef __SOFTMAX_METAX_H__ +#define __SOFTMAX_METAX_H__ + +#include "../softmax.h" + +DESCRIPTOR(metax) + +#endif // __SOFTMAX_METAX_H__ diff --git a/src/infiniop/ops/softmax/metax/softmax_metax.maca b/src/infiniop/ops/softmax/metax/softmax_metax.maca new file mode 100644 index 000000000..3e6b73a56 --- /dev/null +++ b/src/infiniop/ops/softmax/metax/softmax_metax.maca @@ -0,0 +1,120 @@ +#include "../../../devices/metax/metax_common.h" +#include "softmax_metax.h" + +#ifdef ENABLE_METAX_MC_API + #include +#else + #include +#endif +#include "../../../devices/metax/metax_kernel_common.h" + +#include "../../../reduce/cuda/reduce.cuh" + +#include "../cuda/kernel.cuh" + +template +INFINIOP_METAX_KERNEL blockSoftmax( + Tdata *y, const Tdata *x, + size_t dimsize, + ptrdiff_t stride) { + blockSoftmaxKernel(x, y, dimsize, stride); +} + +template +INFINIOP_METAX_KERNEL warpSoftmax( + Tdata *y, const Tdata *x, + size_t othersize, + size_t dimsize, + ptrdiff_t stride) { + warpSoftmaxKernel(x, y, othersize, dimsize, stride); +} + +namespace op::softmax::metax { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + int axis) { + auto info = SoftmaxInfo::create(y_desc, x_desc, axis); + CHECK_RESULT(info); + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +// dimsize > 1024 -> one block per row reduces the full row; otherwise a warp-per-row +// path. Mirrors the NVIDIA softmax launch heuristics. +template +static void dispatchByDimsize(void *y, const void *x, + size_t othersize, size_t dimsize, ptrdiff_t stride, + int num_blocks, hcStream_t stream) { + if (dimsize > 1024) { + blockSoftmax + <<>>((Tdata *)y, (const Tdata *)x, dimsize, stride); + } else if (dimsize > 31) { + constexpr unsigned int BLOCK_SIZE_x = 32; + constexpr unsigned int BLOCK_SIZE_y = 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 + <<>>((Tdata *)y, (const Tdata *)x, othersize, dimsize, stride); + } else { + constexpr unsigned int BLOCK_SIZE_x = 16; + constexpr unsigned int BLOCK_SIZE_y = 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 + <<>>((Tdata *)y, (const Tdata *)x, othersize, dimsize, stride); + } +} + +template +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; + if (dtype == INFINI_DTYPE_F16) { + dispatchByDimsize(y, x, othersize, dimsize, stride, num_blocks, stream); + } else if (dtype == INFINI_DTYPE_BF16) { + dispatchByDimsize<__hpcc_bfloat16, BLOCK_SIZE>(y, x, othersize, dimsize, stride, num_blocks, stream); + } else if (dtype == INFINI_DTYPE_F32) { + dispatchByDimsize(y, x, othersize, dimsize, stride, num_blocks, stream); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + 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( + y, x, _info.dtype, _info.othersize, _info.dimsize, _info.stride, stream)); + } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { + CHECK_STATUS(launchKernel( + y, x, _info.dtype, _info.othersize, _info.dimsize, _info.stride, stream)); + } else { + return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::softmax::metax diff --git a/src/infiniop/ops/softmax/operator.cc b/src/infiniop/ops/softmax/operator.cc index 1664e257f..3c52122ec 100644 --- a/src/infiniop/ops/softmax/operator.cc +++ b/src/infiniop/ops/softmax/operator.cc @@ -9,6 +9,9 @@ #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/softmax_nvidia.cuh" #endif +#ifdef ENABLE_METAX_API +#include "metax/softmax_metax.h" +#endif __INFINI_C infiniStatus_t infiniopCreateSoftmaxDescriptor( infiniopHandle_t handle, @@ -43,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateSoftmaxDescriptor( #endif #ifdef ENABLE_CAMBRICON_API CREATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -74,6 +80,9 @@ __INFINI_C infiniStatus_t infiniopGetSoftmaxWorkspaceSize(infiniopSoftmaxDescrip #endif #ifdef ENABLE_CAMBRICON_API GET(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -110,6 +119,9 @@ __INFINI_C infiniStatus_t infiniopSoftmax( #endif #ifdef ENABLE_CAMBRICON_API CALCULATE(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -141,6 +153,9 @@ __INFINI_C infiniStatus_t infiniopDestroySoftmaxDescriptor(infiniopSoftmaxDescri #endif #ifdef ENABLE_CAMBRICON_API DESTROY(INFINI_DEVICE_CAMBRICON, bang); +#endif +#ifdef ENABLE_METAX_API + DESTROY(INFINI_DEVICE_METAX, metax) #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua index 282f5876f..a69f31a47 100644 --- a/xmake/nvidia.lua +++ b/xmake/nvidia.lua @@ -256,8 +256,8 @@ target("infiniccl-nvidia") if not is_plat("windows") then add_cuflags("-Xcompiler=-fPIC", {force = true}) add_culdflags("-Xcompiler=-fPIC", {force = true}) - add_cxflags("-fPIC") - add_cxxflags("-fPIC") + add_cxflags("-fPIC", {force = true}) + add_cxxflags("-fPIC", {force = true}) local nccl_root = os.getenv("NCCL_ROOT") if nccl_root then