[2026春季][T2-1-1] goog00

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 <cuda_bf16.h>
-#include <cuda_fp16.h>
+#include <cmath>
 
 namespace op::quickgelu::cuda {
 
 typedef struct QuickGeluOp {
 public:
     static constexpr size_t num_inputs = 1;
-
     template <typename T>
     __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<T, half2>) {
-            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<T, cuda_bfloat16>) {
+            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<T, half>) {
-            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<T, __nv_bfloat16>) {
-            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<T, float>) {
-            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<double>(alpha) * x;
-            return x / (1.0 + exp(-ax));
+            return x / (1.0f + expf(-alpha * x));
+        } else {
+            return x / (1.0 + exp(-static_cast<double>(alpha) * x));
         }
     }
-
 } QuickGeluOp;
 
 } // namespace op::quickgelu::cuda

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__

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<infiniopTensorDescriptor_t> input_desc_vec) {
+
+    auto handle = reinterpret_cast<device::metax::Handle *>(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<const void *> 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<T, __nv_bfloat16>. 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

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 {
 

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;

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__

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 <cub/block/block_reduce.cuh>
+#else
+    #include <hccub/block/block_reduce.cuh>
+#endif
+#include "../../../devices/metax/metax_kernel_common.h"
+
+#include "../../../reduce/cuda/reduce.cuh"
+
+#include "../cuda/kernel.cuh"
+
+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 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) {
+    warpSoftmaxKernel<Tdata, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>(x, y, othersize, dimsize, stride);
+}
+
+namespace op::softmax::metax {
+
+struct Descriptor::Opaque {
+    std::shared_ptr<device::metax::Handle::Internal> 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<device::metax::Handle *>(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 <typename Tdata, unsigned int BLOCK_SIZE>
+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, BLOCK_SIZE>
+            <<<num_blocks, BLOCK_SIZE, 0, stream>>>((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, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+            <<<grid_dim, block_dim, 0, stream>>>((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, BLOCK_SIZE_x, BLOCK_SIZE_y, numPerThreadx>
+            <<<grid_dim, block_dim, 0, stream>>>((Tdata *)y, (const Tdata *)x, othersize, dimsize, stride);
+    }
+}
+
+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;
+    if (dtype == INFINI_DTYPE_F16) {
+        dispatchByDimsize<half, BLOCK_SIZE>(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<float, BLOCK_SIZE>(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<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;
+}
+
+} // namespace op::softmax::metax

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;

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